pytorch打印自定义网络的每层的名称

最新推荐文章于 2024-03-24 16:12:13 发布

木里先森

最新推荐文章于 2024-03-24 16:12:13 发布

阅读量5.6k

点赞数 1

分类专栏：深度学习 python

本文链接：https://blog.csdn.net/mr_muli/article/details/100834714

版权

python 同时被 2 个专栏收录

176 篇文章 7 订阅

订阅专栏

深度学习

66 篇文章 0 订阅

订阅专栏

pytorch打印自定义网络的每层的名称

import torch
from torchvision import models
from torchsummary import summary
from resnext_MulTask_clothes import resnext50_elastic


data_class=[8, 7]
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# vgg = models.vgg16().to(device)

model = resnext50_elastic(num_classes=data_class) # 原模型
model = torch.nn.DataParallel(model).cuda() # 并行处理
# 已训练好的模型的pth文件
checkpoint = torch.load('06-resnext50_elastic_checkpoint.pth.tar')
model.load_state_dict(checkpoint['state_dict'], strict=False) # 参数加载

summary(model, (3, 224, 224))

参考连接：https://www.jianshu.com/p/97c626d33924

另：
打印resnet152网络的每层的名称

import torch
from torchvision import models
from torchsummary import summary
from resnet_pretrained import resnet152

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = resnet152() # 原模型
model = torch.nn.DataParallel(model).cuda() # 并行处理
# 已训练好的模型的pth文件
checkpoint = torch.load('resnet152-b121ed2d.pth')

model.load_state_dict(checkpoint, strict=False) # 参数加载

summary(model, (3, 224, 224))