使用torchsummary打印torch模型的每层形状参数

HJC256ZY

于 2020-05-29 20:51:20 发布

阅读量2.7k

点赞数 1

分类专栏： torchsummary Pytorch pyhton 文章标签： pytorch 神经网络

本文链接：https://blog.csdn.net/HJC256ZY/article/details/106432549

版权

pyhton 同时被 3 个专栏收录

19 篇文章 0 订阅

订阅专栏

Pytorch

9 篇文章 0 订阅

订阅专栏

torchsummary

1 篇文章 0 订阅

订阅专栏

torchsummary打印模型的每层形状参数

import torch
import torchvision
from torchsummary import summary          #使用 pip install torchsummary
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
vgg = torchvision.models.vgg16().to(device)
summary(vgg, input_size=(3, 224, 224))