tfm.core.train_utils.try_count_flops

Counts and returns model FLOPs.

tfm.core.train_utils.try_count_flops(
    model: Union[tf.Module, tf.keras.Model],
    inputs_kwargs: Optional[Dict[str, Any]] = None,
    output_path: Optional[str] = None
)

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.

The model's FLOPs.