1. 程式人生 > >[ pytorch ] ——基本使用:(5) 計算模型引數量

[ pytorch ] ——基本使用:(5) 計算模型引數量

################
###   模型定義
# -------------
class MyModel(nn.Module):
    def __init__(self, feat_dim):   # input the dim of output fea-map of Resnet:
        super(MyModel, self).__init__()
        ...
    def forward(self, input):   # input is 2048!
        ...
        return x

net = MyModel()

######################################
params = list(net.parameters())
k = 0
for i in params:
    l = 1
    print("該層的結構:" + str(list(i.size())))
    for j in i.size():
        l *= j
    print("該層引數和:" + str(l))
    k = k + l
print("總引數數量和:" + str(k))
######################################