count FLOPs for operation. Args: op_type (str): the type of operation. input_shapes (dict): the shapes of inputs. attrs (dict): the attributes of the operation. Returns: the total FLOPs of the operation.
(op_type: str, input_shapes: dict, attrs: dict)
| 25 | |
| 26 | |
| 27 | def flops(op_type: str, input_shapes: dict, attrs: dict) -> int: |
| 28 | """ |
| 29 | count FLOPs for operation. |
| 30 | |
| 31 | Args: |
| 32 | op_type (str): the type of operation. |
| 33 | input_shapes (dict): the shapes of inputs. |
| 34 | attrs (dict): the attributes of the operation. |
| 35 | |
| 36 | Returns: |
| 37 | the total FLOPs of the operation. |
| 38 | """ |
| 39 | |
| 40 | if op_type not in _FLOPS_COMPUTE_FUNC_MAP: |
| 41 | return 0 |
| 42 | else: |
| 43 | func = _FLOPS_COMPUTE_FUNC_MAP[op_type] |
| 44 | try: |
| 45 | flops = func(input_shapes, attrs) |
| 46 | except Exception as e: |
| 47 | return 0 |
| 48 | return flops |
| 49 | |
| 50 | |
| 51 | def register_flops(op_type): |