| 147 | } |
| 148 | |
| 149 | Status CalculateFlops(const GraphDef& graph, |
| 150 | const std::vector<InputLayerInfo>& inputs, |
| 151 | Session* session, int64* total_flops, |
| 152 | std::unordered_map<string, int64>* flops_by_op) { |
| 153 | std::unordered_set<string> floppable_ops = { |
| 154 | "Conv2D", "MatMul", "QuantizedConv2D", "QuantizedMatMul", |
| 155 | "DepthwiseConv2dNative"}; |
| 156 | |
| 157 | std::set<string> wanted_shapes; |
| 158 | for (const NodeDef& node : graph.node()) { |
| 159 | if (floppable_ops.count(node.op())) { |
| 160 | for (const string& input : node.input()) { |
| 161 | wanted_shapes.insert(input); |
| 162 | } |
| 163 | wanted_shapes.insert(node.name()); |
| 164 | } |
| 165 | } |
| 166 | std::unordered_map<string, TensorShape> found_shapes; |
| 167 | TF_RETURN_IF_ERROR( |
| 168 | GetOutputShapes(inputs, wanted_shapes, session, &found_shapes)); |
| 169 | |
| 170 | *total_flops = 0; |
| 171 | for (const NodeDef& node : graph.node()) { |
| 172 | if (floppable_ops.count(node.op())) { |
| 173 | int64 current_flops = 0; |
| 174 | // This is a very crude approximation to FLOPs that only looks at a few |
| 175 | // op types that commonly form the bulk of the computation for many |
| 176 | // models. It's included here because getting even an approximate value |
| 177 | // for FLOPs is still very useful for estimating utilization, versus a |
| 178 | // device's theoretical maximum FLOPs/second. |
| 179 | if ((node.op() == "Conv2D") || (node.op() == "QuantizedConv2D")) { |
| 180 | const TensorShape& filter_shape = found_shapes[node.input(1)]; |
| 181 | const TensorShape& output_shape = found_shapes[node.name()]; |
| 182 | int64 filter_height = filter_shape.dim_size(0); |
| 183 | int64 filter_width = filter_shape.dim_size(1); |
| 184 | int64 filter_in_depth = filter_shape.dim_size(2); |
| 185 | int64 output_count = output_shape.num_elements(); |
| 186 | current_flops = |
| 187 | output_count * filter_in_depth * filter_height * filter_width * 2; |
| 188 | } else if ((node.op() == "MatMul") || (node.op() == "QuantizedMatMul")) { |
| 189 | const bool transpose_a = node.attr().at("transpose_a").b(); |
| 190 | const TensorShape& a_shape = found_shapes[node.input(0)]; |
| 191 | const TensorShape& output_shape = found_shapes[node.name()]; |
| 192 | int64 k; |
| 193 | if (transpose_a) { |
| 194 | k = a_shape.dim_size(0); |
| 195 | } else { |
| 196 | k = a_shape.dim_size(1); |
| 197 | } |
| 198 | int64 output_count = output_shape.num_elements(); |
| 199 | current_flops = k * output_count * 2; |
| 200 | } else if (node.op() == "DepthwiseConv2dNative") { |
| 201 | const TensorShape& filter_shape = found_shapes[node.input(1)]; |
| 202 | const TensorShape& output_shape = found_shapes[node.name()]; |
| 203 | int64 filter_height = filter_shape.dim_size(0); |
| 204 | int64 filter_width = filter_shape.dim_size(1); |
| 205 | int64 output_count = output_shape.num_elements(); |
| 206 | current_flops = output_count * filter_height * filter_width * 2; |