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


Posted in Python onMay 12, 2021

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

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

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查看网络参数显存占用量等操作

以上为个人经验,希望能给大家一个参考,也希望大家多多支持三水点靠木。如有错误或未考虑完全的地方,望不吝赐教。

Python 相关文章推荐
教大家使用Python SqlAlchemy
Feb 12 Python
Python中在for循环中嵌套使用if和else语句的技巧
Jun 20 Python
Pycharm导入Python包,模块的图文教程
Jun 13 Python
将python图片转为二进制文本的实例
Jan 24 Python
Python删除n行后的其他行方法
Jan 28 Python
Python实现KNN(K-近邻)算法的示例代码
Mar 05 Python
Python3之手动创建迭代器的实例代码
May 22 Python
Python基础学习之类与实例基本用法与注意事项详解
Jun 17 Python
pytorch中的卷积和池化计算方式详解
Jan 03 Python
python中导入 train_test_split提示错误的解决
Jun 19 Python
python3:excel操作之读取数据并返回字典 + 写入的案例
Sep 01 Python
PyCharm中关于安装第三方包的三个建议
Sep 17 Python
Python入门之使用pandas分析excel数据
May 12 #Python
将Python代码打包成.exe可执行文件的完整步骤
python3实现Dijkstra算法最短路径的实现
pytorch 中autograd.grad()函数的用法说明
python3实现无权最短路径的方法
Python入门之基础语法详解
May 11 #Python
如何利用Matlab制作一款真正的拼图小游戏
You might like
php中数字0和空值的区别分析
2014/06/05 PHP
仿dedecms下拉分页样式修改的thinkphp分页类实例
2014/10/30 PHP
php提交表单时保留多个空格及换行的文本样式的方法
2017/06/20 PHP
清除js缓存的多种方法总结
2016/12/09 Javascript
基于nodejs 的多页面爬虫实例代码
2017/05/31 NodeJs
在 Angular6 中使用 HTTP 请求服务端数据的步骤详解
2018/08/06 Javascript
小程序登录态管理的方法示例
2018/11/13 Javascript
使用express获取微信小程序二维码小记
2019/05/21 Javascript
vue2 中二级路由高亮问题及配置方法
2019/06/10 Javascript
vue路由教程之静态路由
2019/09/03 Javascript
vue中路由跳转不计入history的操作
2020/09/21 Javascript
Vue中引入svg图标的两种方式
2021/01/14 Vue.js
python之yield表达式学习
2014/09/02 Python
Python求两个文本文件以行为单位的交集、并集与差集的方法
2015/06/17 Python
python数据结构之图深度优先和广度优先实例详解
2015/07/08 Python
Python实现字典按照value进行排序的方法分析
2017/12/23 Python
单利模式及python实现方式详解
2018/03/20 Python
python使用turtle绘制分形树
2018/06/22 Python
python-tkinter之按钮的使用,开关方法
2019/06/11 Python
python and or用法详解
2019/06/26 Python
用Python+OpenCV对比图像质量的几种方法
2019/07/15 Python
解决Python计算矩阵乘向量,矩阵乘实数的一些小错误
2019/08/26 Python
xadmin使用formfield_for_dbfield函数过滤下拉表单实例
2020/04/07 Python
TensorFlow Autodiff自动微分详解
2020/07/06 Python
如何利用Python给自己的头像加一个小国旗(小月饼)
2020/10/02 Python
Mio Skincare中文官网:肌肤和身体护理
2016/10/26 全球购物
全球异乡人的跨境社交电商平台:Kouhigh口嗨网
2020/07/24 全球购物
新护士岗前培训制度
2014/02/02 职场文书
培训楼经理岗位责任制
2014/02/10 职场文书
超市中秋节促销方案
2014/03/21 职场文书
安全宣传标语口号
2014/06/06 职场文书
2014大学生党员评议个人总结
2014/09/22 职场文书
2014年城市管理工作总结
2014/12/02 职场文书
2014年药品销售工作总结
2014/12/16 职场文书
小学班主任事迹材料
2014/12/17 职场文书
python实现层次聚类的方法
2021/11/01 Python