最近想到

PyTorch 要如何知道一個模型參數大小

並且顯示層數

類似 keras summary?

原來有 torchsummary

也可以配合 timm 的模型來顯示

 

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.parameter import Parameter
import timm
from torchsummary import summary

class AvgPool(nn.Module):
    def __init__(self):
        super(AvgPool, self).__init__()

    def forward(self, x):
        return torch.mean(x, dim=-1)
# 模型範例
class NModels(nn.Module):
    def __init__(self, modelKey, num_classes=2, pretrained=False):
        super(NModels, self).__init__()
        self.base_model = timm.create_model(modelKey, pretrained=pretrained).cuda()
        
        self.features = nn.Sequential(*list(self.base_model.children())[:-2]).cuda()
        
        self.avgpool = AvgPool().cuda()
        self.fc = nn.Linear(144, num_classes).cuda()
        
    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)
        return x


model = NModels('vit_small_patch32_384')

summary(model, input_size=(3, 384, 384))
total_params = sum(p.numel() for p in model.parameters()) #計算參數量
print("Total Parameters:", total_params)

 

顯示出來的結果相當詳細

 

 


----------------------------------------------------------------
 Layer (type) Output Shape Param #
================================================================
 Conv2d-1 [-1, 384, 12, 12] 1,180,032
 Conv2d-2 [-1, 384, 12, 12] 1,180,032
 Identity-3 [-1, 144, 384] 0
 Identity-4 [-1, 144, 384] 0
......
 Linear-522 [-1, 2] 290
================================================================


 


Total params: 44,949,026
Trainable params: 44,949,026
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 1.69
Forward/backward pass size (MB): 373.78
Params size (MB): 171.47
Estimated Total Size (MB): 546.94
----------------------------------------------------------------
Total Parameters: 22915722

 

 

對於一開始建構模型很有幫助

尤其是需要微調與修改的

給大家參考囉