web-dev-qa-db-ja.com

PyTorchモデルのパラメーターの総数を確認する

PyTorchモデルのパラメーターの総数を数える方法は? Kerasのmodel.count_params()に似たもの。

27
Fábio Perez

PyTorchにはKerasのようにパラメーターの合計数を計算する機能はありませんが、すべてのパラメーターグループの要素数を合計することは可能です。

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

trainableパラメーターのみを計算する場合:

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

これに触発された回答 answer PyTorch Forums

注:私は 自分の質問に答える です。誰かがより良い解決策を持っている場合は、私たちと共有してください。

56
Fábio Perez

モデルをインスタンス化せずに各レイヤーの重みとバイアスの数を計算したい場合、生のファイルをロードして、結果のcollections.OrderedDictを次のように繰り返すことができます:

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)))

次のようなものが得られます

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