web-dev-qa-db-fra.com

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

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

deep-learningpytorch
27
Fábio Perez

PyTorch n'a pas de fonction pour calculer le nombre total de paramètres comme le fait Keras, mais il est possible de faire la somme du 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 les paramètres pouvant être entraînés :

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

Répondez inspiré par ceci réponse sur les forums PyTorch .

Remarque: je suis répondant à ma propre question . Si quelqu'un a une meilleure solution, veuillez partager avec nous.

56
Fábio Perez

Si vous voulez calculer le nombre de poids et de biais dans chaque couche sans instancier le modèle, vous pouvez simplement charger le fichier brut et l'itérer sur le collections.OrderedDict Résultant, comme suit:

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
3
Zhanwen Chen