19 #ifdef INCLUDE_TENSORFLOW 23 #include "allheaders.h" 27 using tensorflow::Status;
28 using tensorflow::Tensor;
29 using tensorflow::TensorShape;
35 TFNetwork::~TFNetwork() {}
37 int TFNetwork::InitFromProtoStr(
const string& proto_str) {
38 if (!model_proto_.ParseFromString(proto_str))
return 0;
39 return InitFromProto();
44 bool TFNetwork::Serialize(TFile* fp)
const {
47 model_proto_.SerializeToString(&proto_str);
50 memcpy(&data[0], proto_str.data(), proto_str.
size());
57 bool TFNetwork::DeSerialize(TFile* fp) {
60 if (!model_proto_.ParseFromArray(&data[0], data.
size())) {
63 return InitFromProto();
68 void TFNetwork::Forward(
bool debug,
const NetworkIO& input,
69 const TransposedArray* input_transpose,
70 NetworkScratch* scratch, NetworkIO* output) {
71 std::vector<std::pair<string, Tensor>> tf_inputs;
72 int depth = input_shape_.depth();
75 const StrideMap& stride_map = input.stride_map();
79 Tensor input_tensor(tensorflow::DT_FLOAT, shape);
81 auto eigen_tensor = input_tensor.flat<
float>();
82 memcpy(eigen_tensor.data(), input.f(0),
83 input.Width() * depth *
sizeof(input.f(0)[0]));
85 tf_inputs.emplace_back(model_proto_.image_input(), input_tensor);
92 if (!model_proto_.image_widths().empty()) {
93 TensorShape size_shape{1};
94 Tensor width_tensor(tensorflow::DT_INT64, size_shape);
95 auto eigen_wtensor = width_tensor.flat<int64>();
96 *eigen_wtensor.data() = stride_map.Size(
FD_WIDTH);
97 tf_inputs.emplace_back(model_proto_.image_widths(), width_tensor);
99 if (!model_proto_.image_heights().empty()) {
100 TensorShape size_shape{1};
101 Tensor height_tensor(tensorflow::DT_INT64, size_shape);
102 auto eigen_htensor = height_tensor.flat<int64>();
103 *eigen_htensor.data() = stride_map.Size(
FD_HEIGHT);
104 tf_inputs.emplace_back(model_proto_.image_heights(), height_tensor);
106 std::vector<string> target_layers = {model_proto_.output_layer()};
107 std::vector<Tensor> outputs;
108 Status s = session_->Run(tf_inputs, target_layers, {}, &outputs);
109 if (!s.ok())
tprintf(
"session->Run failed:%s\n", s.error_message().c_str());
112 const Tensor& output_tensor = outputs[0];
115 int output_batch = output_tensor.shape().dim_size(0);
116 int output_steps = output_tensor.shape().dim_size(1);
117 int output_depth = output_tensor.shape().dim_size(2);
119 ASSERT_HOST(output_depth == output_shape_.depth());
120 output->Resize2d(
false, output_steps, output_depth);
121 auto eigen_output = output_tensor.flat<
float>();
122 memcpy(output->f(0), eigen_output.data(),
123 output_steps * output_depth *
sizeof(output->f(0)[0]));
126 int TFNetwork::InitFromProto() {
127 spec_ = model_proto_.spec();
128 input_shape_.SetShape(
129 model_proto_.batch_size(),
std::max(0, model_proto_.y_size()),
130 std::max(0, model_proto_.x_size()), model_proto_.depth());
131 output_shape_.SetShape(model_proto_.batch_size(), 1, 0,
132 model_proto_.num_classes());
133 output_shape_.set_loss_type(model_proto_.using_ctc() ?
LT_CTC :
LT_SOFTMAX);
134 ni_ = input_shape_.height();
135 no_ = output_shape_.depth();
138 tensorflow::SessionOptions options;
139 session_.reset(NewSession(options));
140 Status s = session_->Create(model_proto_.graph());
141 if (s.ok())
return model_proto_.global_step();
142 tprintf(
"Session_->Create returned '%s'\n", s.error_message().c_str());
148 #endif // ifdef INCLUDE_TENSORFLOW bool DeSerialize(bool swap, FILE *fp)
void resize_no_init(int size)
bool Serialize(FILE *fp) const
virtual bool Serialize(TFile *fp) const