Vérifier le nombre total de paramètres dans un modèle PyTorch

Question

Comment compter le nombre total de paramètres dans un modèle PyTorch ? Quelque chose de similaire à model.count_params() dans Keras.

Answer 1

PyTorch ne dispose pas d'une fonction permettant de calculer le nombre total de paramètres comme le fait Keras, mais il est possible d'additionner le nombre d'éléments pour chaque groupe de paramètres :

pytorch_total_params = sum(p.numel() for p in model.parameters())

Si vous voulez calculer uniquement le formable paramètres :

pytorch_total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

Réponse inspirée par ce qui suit réponse sur les forums PyTorch .

Note : Je suis répondre à ma propre question . Si quelqu'un a une meilleure solution, merci de la partager avec nous.

Answer 2

Pour obtenir le nombre de paramètres de chaque couche comme Keras, PyTorch a modèle.paramètres_nommés() qui renvoie un itérateur à la fois du nom du paramètre et du paramètre lui-même.

Voici un exemple :

from prettytable import PrettyTable

def count_parameters(model):
    table = PrettyTable(["Modules", "Parameters"])
    total_params = 0
    for name, parameter in model.named_parameters():
        if not parameter.requires_grad: continue
        param = parameter.numel()
        table.add_row([name, param])
        total_params+=param
    print(table)
    print(f"Total Trainable Params: {total_params}")
    return total_params

count_parameters(net)

Le résultat ressemblerait à ceci :

+-------------------+------------+
|      Modules      | Parameters |
+-------------------+------------+
| embeddings.weight |   922866   |
|    conv1.weight   |  1048576   |
|     conv1.bias    |    1024    |
|     bn1.weight    |    1024    |
|      bn1.bias     |    1024    |
|    conv2.weight   |  2097152   |
|     conv2.bias    |    1024    |
|     bn2.weight    |    1024    |
|      bn2.bias     |    1024    |
|    conv3.weight   |  2097152   |
|     conv3.bias    |    1024    |
|     bn3.weight    |    1024    |
|      bn3.bias     |    1024    |
|    lin1.weight    |  50331648  |
|     lin1.bias     |    512     |
|    lin2.weight    |   265728   |
|     lin2.bias     |    519     |
+-------------------+------------+
Total Trainable Params: 56773369

Answer 3

Si vous voulez éviter de compter deux fois les paramètres partagés, vous pouvez utiliser torch.Tensor.data_ptr . Par exemple :

sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values())

Voici une implémentation plus verbeuse qui inclut une option pour filtrer les paramètres non entraînables :

def numel(m: torch.nn.Module, only_trainable: bool = False):
    """
    returns the total number of parameters used by `m` (only counting
    shared parameters once); if `only_trainable` is True, then only
    includes parameters with `requires_grad = True`
    """
    parameters = m.parameters()
    if only_trainable:
        parameters = list(p for p in parameters if p.requires_grad)
    unique = dict((p.data_ptr(), p) for p in parameters).values()
    return sum(p.numel() for p in unique)

Answer 4

Vous pouvez utiliser torchsummary pour faire la même chose. C'est juste deux lignes de code.

from torchsummary import summary

print(summary(model, (input_shape)))

Answer 5

Si vous souhaitez calculer le nombre de poids et de biais dans chaque couche sans instancier le modèle, vous pouvez simplement charger le fichier brut et itérer sur le résultat de l'opération. collections.OrderedDict comme ça :

import torch

tensor_dict = torch.load('model.dat', map_location='cpu') # OrderedDict
tensor_list = list(tensor_dict.items())
for layer_tensor_name, tensor in tensor_list:
    print('Layer {}: {} elements'.format(layer_tensor_name, torch.numel(tensor)))

Vous obtiendrez quelque chose comme

conv1.weight: 312
conv1.bias: 26
batch_norm1.weight: 26
batch_norm1.bias: 26
batch_norm1.running_mean: 26
batch_norm1.running_var: 26
conv2.weight: 2340
conv2.bias: 10
batch_norm2.weight: 10
batch_norm2.bias: 10
batch_norm2.running_mean: 10
batch_norm2.running_var: 10
fcs.layers.0.weight: 135200
fcs.layers.0.bias: 260
fcs.layers.1.weight: 33800
fcs.layers.1.bias: 130
fcs.batch_norm_layers.0.weight: 260
fcs.batch_norm_layers.0.bias: 260
fcs.batch_norm_layers.0.running_mean: 260
fcs.batch_norm_layers.0.running_var: 260