Counts and returns model FLOPs. Args: model: A model instance. inputs_kwargs: An optional dictionary of argument pairs specifying inputs' shape specifications to getting corresponding concrete function. output_path: A file path to write the profiling results to. Returns:
(model: Union[tf.Module, tf_keras.Model],
inputs_kwargs: Optional[Dict[str, Any]] = None,
output_path: Optional[str] = None)
| 520 | |
| 521 | |
| 522 | def try_count_flops(model: Union[tf.Module, tf_keras.Model], |
| 523 | inputs_kwargs: Optional[Dict[str, Any]] = None, |
| 524 | output_path: Optional[str] = None): |
| 525 | """Counts and returns model FLOPs. |
| 526 | |
| 527 | Args: |
| 528 | model: A model instance. |
| 529 | inputs_kwargs: An optional dictionary of argument pairs specifying inputs' |
| 530 | shape specifications to getting corresponding concrete function. |
| 531 | output_path: A file path to write the profiling results to. |
| 532 | |
| 533 | Returns: |
| 534 | The model's FLOPs. |
| 535 | """ |
| 536 | if hasattr(model, 'inputs'): |
| 537 | try: |
| 538 | # Get input shape and set batch size to 1. |
| 539 | if model.inputs: |
| 540 | inputs = [ |
| 541 | tf.TensorSpec([1] + input.shape[1:], input.dtype) |
| 542 | for input in model.inputs |
| 543 | ] |
| 544 | concrete_func = tf.function(model).get_concrete_function(inputs) |
| 545 | # If model.inputs is invalid, try to use the input to get concrete |
| 546 | # function for model.call (subclass model). |
| 547 | else: |
| 548 | concrete_func = tf.function(model.call).get_concrete_function( |
| 549 | **inputs_kwargs) |
| 550 | frozen_func, _ = convert_variables_to_constants_v2_as_graph(concrete_func) |
| 551 | |
| 552 | # Calculate FLOPs. |
| 553 | run_meta = tf.compat.v1.RunMetadata() |
| 554 | opts = tf.compat.v1.profiler.ProfileOptionBuilder.float_operation() |
| 555 | if output_path is not None: |
| 556 | opts['output'] = f'file:outfile={output_path}' |
| 557 | else: |
| 558 | opts['output'] = 'none' |
| 559 | flops = tf.compat.v1.profiler.profile( |
| 560 | graph=frozen_func.graph, run_meta=run_meta, options=opts) |
| 561 | return flops.total_float_ops |
| 562 | except Exception as e: # pylint: disable=broad-except |
| 563 | logging.info( |
| 564 | 'Failed to count model FLOPs with error %s, because the build() ' |
| 565 | 'methods in keras layers were not called. This is probably because ' |
| 566 | 'the model was not feed any input, e.g., the max train step already ' |
| 567 | 'reached before this run.', e) |
| 568 | return None |
| 569 | return None |
| 570 | |
| 571 | |
| 572 | @ops.RegisterStatistics('Einsum', 'flops') |