App下載

pytorch查看網(wǎng)絡參數(shù)顯存占用量等操作

來源: 一語呢喃醉 2021-08-18 14:39:09 瀏覽數(shù) (7349)
反饋

pytorch是一種比較吃顯存的python應用。有些時候我們調(diào)試代碼需要關注顯存占用量,在調(diào)試代碼的時候也有查看網(wǎng)絡參數(shù)的需要。這些調(diào)試是需要的參數(shù)我們在調(diào)試工具中就能看到,那么有哪些調(diào)試工具可以看這些參數(shù)呢?又該怎么看這些參數(shù)呢?接下來的這篇文章帶你了解。

1.使用torchstat

pip install torchstat 

from torchstat import stat
import torchvision.models as models
model = models.resnet152()
stat(model, (3, 224, 224))

關于stat函數(shù)的參數(shù),第一個應該是模型,第二個則是輸入尺寸,3為通道數(shù)。我沒有調(diào)研該函數(shù)的詳細參數(shù),也不知道為什么使用的時候并不提示相應的參數(shù)。

2.使用torchsummary

pip install torchsummary
 
from torchsummary import summary
summary(model.cuda(),input_size=(3,32,32),batch_size=-1)

使用該函數(shù)直接對參數(shù)進行提示,可以發(fā)現(xiàn)直接有顯式輸入batch_size的地方,我自己的感覺好像該函數(shù)更好一些。但是?。。〔恢罏槭裁?,該函數(shù)在我的機器上一直報錯!??!

TypeError: can't convert CUDA tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.

Update:經(jīng)過論壇咨詢,報錯的原因找到了,只需要把

pip install torchsummary

修改為

pip install torch-summary

補充:Pytorch查看模型參數(shù)并計算模型參數(shù)量與可訓練參數(shù)量

查看模型參數(shù)(以AlexNet為例)

import torch
import torch.nn as nn
import torchvision
class AlexNet(nn.Module):
    def __init__(self,num_classes=1000):
        super(AlexNet,self).__init__()
        self.feature_extraction = nn.Sequential(
            nn.Conv2d(in_channels=3,out_channels=96,kernel_size=11,stride=4,padding=2,bias=False),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3,stride=2,padding=0),
            nn.Conv2d(in_channels=96,out_channels=192,kernel_size=5,stride=1,padding=2,bias=False),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3,stride=2,padding=0),
            nn.Conv2d(in_channels=192,out_channels=384,kernel_size=3,stride=1,padding=1,bias=False),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=384,out_channels=256,kernel_size=3,stride=1,padding=1,bias=False),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=256,out_channels=256,kernel_size=3,stride=1,padding=1,bias=False),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=0),
        )
        self.classifier = nn.Sequential(
            nn.Dropout(p=0.5),
            nn.Linear(in_features=256*6*6,out_features=4096),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.5),
            nn.Linear(in_features=4096, out_features=4096),
            nn.ReLU(inplace=True),
            nn.Linear(in_features=4096, out_features=num_classes),
        )
    def forward(self,x):
        x = self.feature_extraction(x)
        x = x.view(x.size(0),256*6*6)
        x = self.classifier(x)
        return x
if __name__ =='__main__':
    # model = torchvision.models.AlexNet()
    model = AlexNet()
    
    # 打印模型參數(shù)
    #for param in model.parameters():
        #print(param)
    
    #打印模型名稱與shape
    for name,parameters in model.named_parameters():
        print(name,':',parameters.size())
feature_extraction.0.weight : torch.Size([96, 3, 11, 11])
feature_extraction.3.weight : torch.Size([192, 96, 5, 5])
feature_extraction.6.weight : torch.Size([384, 192, 3, 3])
feature_extraction.8.weight : torch.Size([256, 384, 3, 3])
feature_extraction.10.weight : torch.Size([256, 256, 3, 3])
classifier.1.weight : torch.Size([4096, 9216])
classifier.1.bias : torch.Size([4096])
classifier.4.weight : torch.Size([4096, 4096])
classifier.4.bias : torch.Size([4096])
classifier.6.weight : torch.Size([1000, 4096])
classifier.6.bias : torch.Size([1000])

計算參數(shù)量與可訓練參數(shù)量

def get_parameter_number(model):
    total_num = sum(p.numel() for p in model.parameters())
    trainable_num = sum(p.numel() for p in model.parameters() if p.requires_grad)
    return {'Total': total_num, 'Trainable': trainable_num}

第三方工具

from torchstat import stat
import torchvision.models as models
model = models.alexnet()
stat(model, (3, 224, 224))

在這里插入圖片描述

from torchvision.models import alexnet
import torch
from thop import profile
model = alexnet()
input = torch.randn(1, 3, 224, 224)
flops, params = profile(model, inputs=(input, ))
print(flops, params)

在這里插入圖片描述

以上就是pytorch查看網(wǎng)絡參數(shù)顯存占用量等操作的全部內(nèi)容,希望能給大家一個參考,也希望大家多多支持W3Cschool。如有錯誤或未考慮完全的地方,望不吝賜教。


0 人點贊