Acquisition des paramètres du modèle (Parameters) et de la quantité de calcul (FLOPs) [à l'aide de thop]

Conseils : Pour certains codes open source, il n'y a pas de code pour calculer les paramètres et les calculs du réseau. Voici une méthode générale pour obtenir la quantité de paramètre et la quantité de calcul du réseau. Utilisez thop pour l'obtenir rapidement

1 Paramètres du modèle et calcul

Nombre de paramètres #paramsC'est-à-dire combien de paramètres sont inclus dans le modèle de réseau, ce qui n'a rien à voir avec les données d'entrée, mais a principalement une relation avec la structure du modèle ; sa principale influence sur le fonctionnement du modèle est la mémoire ou la mémoire vidéo requise

Montant de calcul #FLOPsHabituellement, les FLOP (opérations en virgule flottante, le nombre d'opérations en virgule flottante) sont utilisés pour représenter la quantité de calcul, qui mesure principalement la complexité de l'algorithme/modèle. Dans l'article, il est généralement représenté par des GFLOP, 1GFLOP = 10^9 FLOP ;

2 étapes d'installation

pip install thop

3 exemples de code

#### (tel que le calcul du montant du paramètre et le montant du calcul du réseau suivant) ####

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)

#### (pas de code, vous pouvez l'optimiser par vous-même) ####

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 Affichage des résultats