mmengine.analysis.get_model_complexity_info

mmengine.analysis.get_model_complexity_info(model, input_shape=None, inputs=None, show_table=True, show_arch=True)

Interface to get the complexity of a model.

The parameter inputs are fed to the forward method of model. If inputs is not specified, the input_shape is required and it will be used to construct the dummy input fed to model. If the forward of model requires two or more inputs, the inputs should be a tuple of tensor or the input_shape should be a tuple of tuple which each element will be constructed into a dumpy input.

实际案例

>>> # the forward of model accepts only one input
>>> input_shape = (3, 224, 224)
>>> get_model_complexity_info(model, input_shape=input_shape)
>>> # the forward of model accepts two or more inputs
>>> input_shape = ((3, 224, 224), (3, 10))
>>> get_model_complexity_info(model, input_shape=input_shape)

参数

• model (nn.Module) – The model to analyze.

• input_shape (Union[Tuple[int, ...], Tuple[Tuple[int, ...]], None]) – The input shape of the model. If “inputs” is not specified, the “input_shape” should be set. Defaults to None.

• inputs (torch.Tensor, tuple[torch.Tensor, ...] or Tuple[Any, ...], optional]) – The input tensor(s) of the model. If not given the input tensor will be generated automatically with the given input_shape. Defaults to None.

• show_table (bool) – Whether to show the complexity table. Defaults to True.

• show_arch (bool) – Whether to show the complexity arch. Defaults to True.

返回

The complexity information of the model.

返回类型

dict