02. 深度学习网络结构可视化

1. torchsummary

安装:

pip install torchsummary

使用方法：

import torch
from torchvision import models
from torchsummary import summary

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
vgg = models.vgg16().to(device)

'''
# summary 函数完整形式
summary(model, input_size, batch_size=-1, device=torch.device('cuda:0'),dtypes=None)
'''
summary(vgg, (3, 224, 224))

输出大概长这个样子：

官方链接及使用指南：https://github.com/sksq96/pytorch-summary

Post Author: jasonyang