Adquisición de parámetros del modelo (Parameters) y cantidad de cálculo (FLOP) [usando thop]

Sugerencias: para algunos códigos fuente abiertos, no hay ningún código para calcular los parámetros y cálculos de la red. Aquí hay un método general para obtener la cantidad de parámetros y la cantidad de cálculo de la red. Usa thop para obtenerlo rápidamente.

1 Parámetros del modelo y cálculo

Cantidad de parámetros #paramsEs decir, cuantos parámetros se incluyen en el modelo de red, que no tiene nada que ver con los datos de entrada, sino que principalmente tiene relación con la estructura del modelo, su principal influencia en el funcionamiento del modelo es la memoria requerida o memoria de video.

Cantidad de cálculo #FLOPPor lo general, los FLOP (operaciones de punto flotante, el número de operaciones de punto flotante) se utilizan para representar la cantidad de cálculo, que mide principalmente la complejidad del algoritmo/modelo. En el documento, generalmente se representa mediante GFLOP, 1GFLOP = 10^9 FLOP;

2 tienda de instalación

pip install thop

3 código de muestra

#### (como calcular la cantidad de parámetro y la cantidad de cálculo de la siguiente red) ####

import torch
import torch.nn as nn

# SENet
class SELayer(nn.Module):
    def __init__(self, channel, reduction=16):
        super(SELayer, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(channel, channel // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(channel // reduction, channel, bias=False),
            nn.Sigmoid()
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y.expand_as(x)

#### (código superior, puede optimizarlo usted mismo) ####

import thop
if __name__ == '__main__':
    # 输入1 channel
    model = SELayer(channel=64)
    # (1, 64, 640, 640) 输入的图片尺寸
    x = torch.randn(1, 64, 640, 640)
    flops, params = thop.profile(model, inputs=(x,))
    print("%s | %s | %s" % ("Model", "Params(M)", "FLOPs(G)"))
    print("------|-----------|------")
    print("%s | %.7f | %.7f" % ("模型  ", params / (1000 ** 2), flops / (1000 ** 3)))

4 Visualización de resultados