18 #ifdef INCLUDE_TENSORFLOW
22 #include "allheaders.h"
26 using tensorflow::Status;
27 using tensorflow::Tensor;
28 using tensorflow::TensorShape;
34 int TFNetwork::InitFromProtoStr(
const std::string& proto_str) {
35 if (!model_proto_.ParseFromString(proto_str))
return 0;
36 return InitFromProto();
44 model_proto_.SerializeToString(&proto_str);
47 memcpy(&data[0], proto_str.data(), proto_str.size());
57 if (!model_proto_.ParseFromArray(&data[0], data.
size())) {
60 return InitFromProto();
65 void TFNetwork::Forward(
bool debug,
const NetworkIO& input,
66 const TransposedArray* input_transpose,
67 NetworkScratch* scratch, NetworkIO* output) {
68 std::vector<std::pair<std::string, Tensor>> tf_inputs;
69 int depth = input_shape_.depth();
72 const StrideMap& stride_map = input.stride_map();
76 Tensor input_tensor(tensorflow::DT_FLOAT, shape);
78 auto eigen_tensor = input_tensor.flat<
float>();
79 memcpy(eigen_tensor.data(), input.f(0),
80 input.Width() * depth *
sizeof(input.f(0)[0]));
82 tf_inputs.emplace_back(model_proto_.image_input(), input_tensor);
89 if (!model_proto_.image_widths().empty()) {
90 TensorShape size_shape{1};
91 Tensor width_tensor(tensorflow::DT_INT64, size_shape);
92 auto eigen_wtensor = width_tensor.flat<tensorflow::int64>();
93 *eigen_wtensor.data() = stride_map.Size(
FD_WIDTH);
94 tf_inputs.emplace_back(model_proto_.image_widths(), width_tensor);
96 if (!model_proto_.image_heights().empty()) {
97 TensorShape size_shape{1};
98 Tensor height_tensor(tensorflow::DT_INT64, size_shape);
99 auto eigen_htensor = height_tensor.flat<tensorflow::int64>();
100 *eigen_htensor.data() = stride_map.Size(
FD_HEIGHT);
101 tf_inputs.emplace_back(model_proto_.image_heights(), height_tensor);
103 std::vector<std::string> target_layers = {model_proto_.output_layer()};
104 std::vector<Tensor> outputs;
105 Status s = session_->Run(tf_inputs, target_layers, {}, &outputs);
106 if (!s.ok())
tprintf(
"session->Run failed:%s\n", s.error_message().c_str());
109 const Tensor& output_tensor = outputs[0];
112 int output_batch = output_tensor.shape().dim_size(0);
113 int output_steps = output_tensor.shape().dim_size(1);
114 int output_depth = output_tensor.shape().dim_size(2);
116 ASSERT_HOST(output_depth == output_shape_.depth());
117 output->Resize2d(
false, output_steps, output_depth);
118 auto eigen_output = output_tensor.flat<
float>();
119 memcpy(output->f(0), eigen_output.data(),
120 output_steps * output_depth *
sizeof(output->f(0)[0]));
123 int TFNetwork::InitFromProto() {
124 spec_ = model_proto_.spec();
125 input_shape_.SetShape(
126 model_proto_.batch_size(), std::max(0, model_proto_.y_size()),
127 std::max(0, model_proto_.x_size()), model_proto_.depth());
128 output_shape_.SetShape(model_proto_.batch_size(), 1, 0,
129 model_proto_.num_classes());
130 output_shape_.set_loss_type(model_proto_.using_ctc() ?
LT_CTC :
LT_SOFTMAX);
131 ni_ = input_shape_.height();
132 no_ = output_shape_.depth();
135 tensorflow::SessionOptions options;
136 session_.reset(NewSession(options));
137 Status s = session_->Create(model_proto_.graph());
138 if (s.ok())
return model_proto_.global_step();
139 tprintf(
"Session_->Create returned '%s'\n", s.error_message().c_str());
145 #endif // ifdef INCLUDE_TENSORFLOW