Pytorch高效展示网络结构与参数方法
发布时间
阅读量:
阅读量
首先,需要完成torchsummary的安装步骤。
pip isntall torchsummary
以下为示例代码,其中(3, 100, 100)分别对应单张图像的通道数量、高度以及宽度
from torchsummary import summary
import torchvision
import torch
def model_info(model):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
backbone = model.to(device)
summary(backbone, (3, 100, 100))
model = torchvision.models.vgg16(pretrained=False)
model_info(model)
最终输出结果包含13个卷积层(Conv2d)与3个全连接层
全部评论 (0)
还没有任何评论哟~
