pytorch查看网络参数显存占用量等操作

更新时间：2021年05月12日 11:11:54 作者：张林克

这篇文章主要介绍了pytorch查看网络参数显存占用量等操作，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧

1.使用torchstat

pip install torchstat 

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

关于stat函数的参数，第一个应该是模型，第二个则是输入尺寸，3为通道数。我没有调研该函数的详细参数，也不知道为什么使用的时候并不提示相应的参数。

2.使用torchsummary

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

使用该函数直接对参数进行提示，可以发现直接有显式输入batch_size的地方，我自己的感觉好像该函数更好一些。但是！！！不知道为什么，该函数在我的机器上一直报错！！！

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

Update：经过论坛咨询，报错的原因找到了，只需要把

pip install torchsummary

修改为

pip install torch-summary

补充：Pytorch查看模型参数并计算模型参数量与可训练参数量

查看模型参数(以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()
    
    # 打印模型参数
    #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])

计算参数量与可训练参数量

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)

在这里插入图片描述

以上为个人经验，希望能给大家一个参考，也希望大家多多支持脚本之家。如有错误或未考虑完全的地方，望不吝赐教。

您可能感兴趣的文章:

python 多线程与多进程效率测试
这篇文章主要介绍了python 多线程与多进程效率测试，在Python中，计算密集型任务适用于多进程，IO密集型任务适用于多线程、接下来看看文章得实例吧,需要的朋友可以参考一下哟
2021-10-10
Python tracemalloc跟踪内存分配问题
这篇文章主要介绍了Python tracemalloc跟踪内存分配问题,具有很好的参考价值,希望对大家有所帮助,如有错误或未考虑完全的地方,望不吝赐教
2023-11-11
浅谈Python xlwings 读取Excel文件的正确姿势
这篇文章主要介绍了浅谈Python xlwings 读取Excel文件的正确姿势，文中通过示例代码介绍的非常详细，对大家的学习或者工作具有一定的参考学习价值，需要的朋友们下面随着小编来一起学习学习吧
2021-02-02
使用Python实现嵌套绘图并为条形图添加自定义标注
论文绘图时经常需要多图嵌套,正好最近绘图用到了,所以这篇文章主要为大家详细介绍了如何使用Python实现嵌套绘图并为条形图添加自定义标注,感兴趣的可以了解下
2024-02-02
python使用jenkins发送企业微信通知的实现
公司使用的是企业微信，因此考虑Jenkins通知企业微信机器人的实现方式，本文主要介绍了python使用jenkins发送企业微信通知的实现，感兴趣的可以了解一下
2021-06-06
Python结合OpenCV和Pyzbar实现实时摄像头识别二维码
这篇文章主要为大家详细介绍了如何使用Python编程语言结合OpenCV和Pyzbar库来实时摄像头识别二维码,文中的示例代码讲解详细,需要的可以参考下
2024-01-01
Python使用requests发送POST请求实例代码
这篇文章主要介绍了Python使用requests发送POST请求实例代码，小编觉得还是挺不错的，具有一定借鉴价值，需要的朋友可以参考下
2018-01-01
Django实现聊天机器人
本文基于channels + websocket结合Celery和Python爬虫技术打造了一个会算术懂诗文的聊天机器人，是非常难得的一个Django综合应用项目哦，感兴趣的朋友可以参考下
2021-05-05
Python绘制全球疫情变化地图的实例代码
这篇文章主要介绍了使用Python绘制全球疫情变化地图的实例代码，本文通过实例代码给大家介绍的非常详细，对大家的学习或工作具有一定的参考借鉴价值,需要的朋友可以参考下
2020-04-04
Python 中 f-Strings 的作用
这篇文章主要介绍了Python 中 f-Strings 的作用, f-strings 是用来非常方便的格式化输出的,觉得它的使用方法无外乎就是 print(f'value = { value }'，其实，f-strings 远超你的预期，今天来梳理一下它还能做那些很酷的事情
2021-10-10