mmcv.cnn.get_model_complexity_info

mmcv.cnn.get_model_complexity_info(model: torch.nn.modules.module.Module, input_shape: tuple, print_per_layer_stat: bool = True, as_strings: bool = True, input_constructor: Optional[Callable] = None, flush: bool = False, ost: TextIO = <_io.TextIOWrapper name='<stdout>' mode='w' encoding='UTF-8'>) → tuple

Get complexity information of a model.

This method can calculate FLOPs and parameter counts of a model with corresponding input shape. It can also print complexity information for each layer in a model.

Supported layers are listed as below:

• Convolutions: nn.Conv1d, nn.Conv2d, nn.Conv3d.

• Activations: nn.ReLU, nn.PReLU, nn.ELU, nn.LeakyReLU, nn.ReLU6.

• Poolings: nn.MaxPool1d, nn.MaxPool2d, nn.MaxPool3d, nn.AvgPool1d, nn.AvgPool2d, nn.AvgPool3d, nn.AdaptiveMaxPool1d, nn.AdaptiveMaxPool2d, nn.AdaptiveMaxPool3d, nn.AdaptiveAvgPool1d, nn.AdaptiveAvgPool2d, nn.AdaptiveAvgPool3d.

• BatchNorms: nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.GroupNorm, nn.InstanceNorm1d, InstanceNorm2d, InstanceNorm3d, nn.LayerNorm.

• Linear: nn.Linear.

• Deconvolution: nn.ConvTranspose2d.

• Upsample: nn.Upsample.

Parameters

• model (nn.Module) – The model for complexity calculation.

• input_shape (tuple) – Input shape used for calculation.

• print_per_layer_stat (bool) – Whether to print complexity information for each layer in a model. Default: True.

• as_strings (bool) – Output FLOPs and params counts in a string form. Default: True.

• input_constructor (None | callable) – If specified, it takes a callable method that generates input. otherwise, it will generate a random tensor with input shape to calculate FLOPs. Default: None.

• flush (bool) – same as that in print(). Default: False.

• ost (stream) – same as file param in print(). Default: sys.stdout.

Returns

If as_strings is set to True, it will return FLOPs and parameter counts in a string format. otherwise, it will return those in a float number format.

Return type

tuple[float | str]