| 115 | } |
| 116 | |
| 117 | Status GetOutputShapes(const std::vector<InputLayerInfo>& inputs, |
| 118 | const std::set<string>& wanted_shapes, Session* session, |
| 119 | std::unordered_map<string, TensorShape>* node_shapes) { |
| 120 | std::vector<std::pair<string, tensorflow::Tensor> > input_tensors; |
| 121 | CreateTensorsFromInputInfo(inputs, &input_tensors); |
| 122 | std::vector<tensorflow::Tensor> output_tensors; |
| 123 | std::vector<string> output_tensor_names; |
| 124 | for (const string& wanted_shape : wanted_shapes) { |
| 125 | bool is_input = false; |
| 126 | for (const std::pair<string, tensorflow::Tensor>& input_tensor : |
| 127 | input_tensors) { |
| 128 | if (input_tensor.first == wanted_shape) { |
| 129 | (*node_shapes)[wanted_shape] = input_tensor.second.shape(); |
| 130 | is_input = true; |
| 131 | break; |
| 132 | } |
| 133 | } |
| 134 | if (!is_input) { |
| 135 | output_tensor_names.push_back(wanted_shape); |
| 136 | } |
| 137 | } |
| 138 | TF_RETURN_IF_ERROR( |
| 139 | session->Run(input_tensors, output_tensor_names, {}, &output_tensors)); |
| 140 | CHECK_EQ(output_tensors.size(), output_tensor_names.size()); |
| 141 | for (int i = 0; i < output_tensor_names.size(); ++i) { |
| 142 | const string& wanted_shape_name = output_tensor_names[i]; |
| 143 | const TensorShape& found_shape = output_tensors[i].shape(); |
| 144 | (*node_shapes)[wanted_shape_name] = found_shape; |
| 145 | } |
| 146 | return Status::OK(); |
| 147 | } |
| 148 | |
| 149 | Status CalculateFlops(const GraphDef& graph, |
| 150 | const std::vector<InputLayerInfo>& inputs, |
no test coverage detected