# 模型复杂度分析

我们提供了一个工具来帮助分析网络的复杂性。我们借鉴了 [fvcore](https://github.com/facebookresearch/fvcore) 的实现思路来构建这个工具，并计划在未来支持更多的自定义算子。目前的工具提供了用于计算给定模型的浮点运算量（FLOPs）、激活量（Activations）和参数量（Parameters）的接口，并支持以网络结构或表格的形式逐层打印相关信息，同时提供了算子级别（operator）和模块级别（Module）的统计。如果您对统计浮点运算量的实现细节感兴趣，请参考 [Flop Count](https://github.com/facebookresearch/fvcore/blob/main/docs/flop_count.md)。

## 定义

模型复杂度有 3 个指标，分别是浮点运算量（FLOPs）、激活量（Activations）以及参数量（Parameters），它们的定义如下：

- 浮点运算量

  浮点运算量不是一个定义非常明确的指标，在这里参考 [detectron2](https://detectron2.readthedocs.io/en/latest/modules/fvcore.html#fvcore.nn.FlopCountAnalysis) 的描述，将一组乘加运算定义为 1 个 flop。

- 激活量

  激活量用于衡量某一层产生的特征数量。

- 参数量

  模型的参数量。

例如，给定输入尺寸 `inputs = torch.randn((1, 3, 10, 10))`，和一个卷积层 `conv = nn.Conv2d(in_channels=3, out_channels=10, kernel_size=3)`，那么它输出的特征图尺寸为 `(1, 10, 8, 8)`，则它的浮点运算量是 `17280 = 10*8*8*3*3*3`（10*8*8 表示输出的特征图大小、3*3*3 表示每一个输出需要的计算量）、激活量是 `640 = 10*8*8`、参数量是 `280 = 3*10*3*3 + 10`（3*10*3\*3 表示权重的尺寸、10 表示偏置值的尺寸）。

## 用法

### 基于 `nn.Module` 构建的模型

构建模型

```python
from torch import nn

from mmengine.analysis import get_model_complexity_info


# 以字典的形式返回分析结果，包括:
# ['flops', 'flops_str', 'activations', 'activations_str', 'params', 'params_str', 'out_table', 'out_arch']
class InnerNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 10)
        self.fc2 = nn.Linear(10, 10)

    def forward(self, x):
        return self.fc1(self.fc2(x))


class TestNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 10)
        self.fc2 = nn.Linear(10, 10)
        self.inner = InnerNet()

    def forward(self, x):
        return self.fc1(self.fc2(self.inner(x)))


input_shape = (1, 10)
model = TestNet()
```

`get_model_complexity_info` 返回的 `analysis_results` 是一个包含 7 个值的字典:

- `flops`: flop 的总数, 例如, 1000, 1000000
- `flops_str`: 格式化的字符串, 例如, 1.0G, 1.0M
- `params`: 全部参数的数量, 例如, 1000, 1000000
- `params_str`: 格式化的字符串, 例如, 1.0K, 1M
- `activations`: 激活量的总数, 例如, 1000, 1000000
- `activations_str`: 格式化的字符串, 例如, 1.0G, 1M
- `out_table`: 以表格形式打印相关信息

打印结果

- 以表格形式打印相关信息

  ```python
  print(analysis_results['out_table'])
  ```

  ```text
  +---------------------+----------------------+--------+--------------+
  | module              | #parameters or shape | #flops | #activations |
  +---------------------+----------------------+--------+--------------+
  | model               | 0.44K                | 0.4K   | 40           |
  |  fc1                |  0.11K               |  100   |  10          |
  |   fc1.weight        |   (10, 10)           |        |              |
  |   fc1.bias          |   (10,)              |        |              |
  |  fc2                |  0.11K               |  100   |  10          |
  |   fc2.weight        |   (10, 10)           |        |              |
  |   fc2.bias          |   (10,)              |        |              |
  |  inner              |  0.22K               |  0.2K  |  20          |
  |   inner.fc1         |   0.11K              |   100  |   10         |
  |    inner.fc1.weight |    (10, 10)          |        |              |
  |    inner.fc1.bias   |    (10,)             |        |              |
  |   inner.fc2         |   0.11K              |   100  |   10         |
  |    inner.fc2.weight |    (10, 10)          |        |              |
  |    inner.fc2.bias   |    (10,)             |        |              |
  +---------------------+----------------------+--------+--------------+
  ```

- 以网络层级结构打印相关信息

  ```python
  print(analysis_results['out_arch'])
  ```

  ```bash
  TestNet(
    #params: 0.44K, #flops: 0.4K, #acts: 40
    (fc1): Linear(
      in_features=10, out_features=10, bias=True
      #params: 0.11K, #flops: 100, #acts: 10
    )
    (fc2): Linear(
      in_features=10, out_features=10, bias=True
      #params: 0.11K, #flops: 100, #acts: 10
    )
    (inner): InnerNet(
      #params: 0.22K, #flops: 0.2K, #acts: 20
      (fc1): Linear(
        in_features=10, out_features=10, bias=True
        #params: 0.11K, #flops: 100, #acts: 10
      )
      (fc2): Linear(
        in_features=10, out_features=10, bias=True
        #params: 0.11K, #flops: 100, #acts: 10
      )
    )
  )
  ```

- 以字符串的形式打印结果

  ```python
  print("Model Flops:{}".format(analysis_results['flops_str']))
  # Model Flops:0.4K
  print("Model Parameters:{}".format(analysis_results['params_str']))
  # Model Parameters:0.44K
  ```

### 基于 BaseModel（来自 MMEngine）构建的模型

```python
import torch.nn.functional as F
import torchvision
from mmengine.model import BaseModel
from mmengine.analysis import get_model_complexity_info


class MMResNet50(BaseModel):
    def __init__(self):
        super().__init__()
        self.resnet = torchvision.models.resnet50()

    def forward(self, imgs, labels=None, mode='tensor'):
        x = self.resnet(imgs)
        if mode == 'loss':
            return {'loss': F.cross_entropy(x, labels)}
        elif mode == 'predict':
            return x, labels
        elif mode == 'tensor':
            return x


input_shape = (3, 224, 224)
model = MMResNet50()

analysis_results = get_model_complexity_info(model, input_shape)

print("Model Flops:{}".format(analysis_results['flops_str']))
# Model Flops:4.145G
print("Model Parameters:{}".format(analysis_results['params_str']))
# Model Parameters:25.557M
```

## 其他接口

除了上述基本用法，`get_model_complexity_info` 还能接受以下参数，输出定制化的统计结果：

- `model`: (nn.Module) 待分析的模型
- `input_shape`: (tuple) 输入尺寸，例如 (3, 224, 224)
- `inputs`: (optional: torch.Tensor), 如果传入该参数, `input_shape` 会被忽略
- `show_table`: (bool) 是否以表格形式返回统计结果，默认值：True
- `show_arch`: (bool) 是否以网络结构形式返回统计结果，默认值：True