(graph, detail=False)
| 173 | |
| 174 | |
| 175 | def _graph_flops(graph, detail=False): |
| 176 | assert isinstance(graph, GraphWrapper) |
| 177 | flops = 0 |
| 178 | op_flops = 0 |
| 179 | table = Table(["OP Type", 'Param name', "Flops"]) |
| 180 | for op in graph.ops(): |
| 181 | param_name = '' |
| 182 | if op.type() in ['conv2d', 'depthwise_conv2d']: |
| 183 | op_flops = count_convNd(op) |
| 184 | flops += op_flops |
| 185 | param_name = op.inputs("Filter")[0].name() |
| 186 | elif op.type() == 'pool2d': |
| 187 | op_flops = count_pool2d(op) |
| 188 | flops += op_flops |
| 189 | |
| 190 | elif op.type() in ['mul', 'matmul']: |
| 191 | op_flops = count_linear(op) |
| 192 | flops += op_flops |
| 193 | param_name = op.inputs("Y")[0].name() |
| 194 | elif op.type() == 'batch_norm': |
| 195 | op_flops = count_bn(op) |
| 196 | flops += op_flops |
| 197 | elif op.type().startswith('element'): |
| 198 | op_flops = count_element_op(op) |
| 199 | flops += op_flops |
| 200 | if op_flops != 0: |
| 201 | table.add_row([op.type(), param_name, op_flops]) |
| 202 | op_flops = 0 |
| 203 | if detail: |
| 204 | table.print_table() |
| 205 | return flops |
| 206 | |
| 207 | |
| 208 | def static_flops(program, print_detail=False): |
no test coverage detected