torch加载模型

torch保存的模型有很多形式， .pth.tar 后缀的文件，直接使用 torch.load 直接加载模型就可以了。
1、查看 .pth.tar 文件模型参数

import torch
import torchvision.models as models

checkpoint = torch.load('resnet50_train_60_epochs-c8e5653e.pth.tar')	# 加载模型
print(checkpoint.keys())												# 查看模型元素
state_dict = checkpoint['state_dict']

print(checkpoint['epoch'])
print(checkpoint['arch'])
print(checkpoint['best_prec1'])
# print(checkpoint['ce_optimizer'])
print(state_dict.keys())
print(state_dict['module.bn1.bias'])									# 打印 module.bn1.bias 的权值
print(state_dict['module.bn1.bias'])									# 打印 module.bn1.bias 的形状

2、由 .pth.tar 文件加载模型

import torch
import torchvision.models as models

checkpoint = torch.load('resnet50_train_60_epochs-c8e5653e.pth.tar')
arch = checkpoint['arch']
model = models.__dict__[arch]()
model = torch.nn.DataParallel(model).cuda()
model.load_state_dict(checkpoint['state_dict'])

# print(model.state_dict())
# print(model)

3、保存模型（只保存参数）

只保存模型权重参数，不保存模型结构。

torch.save(model.state_dict(), 'resnet_only_paramater.pth')

4、保存模型（保存全部）

保存整个模型，包括模型权重和模型结构。

torch.save(model, 'resnet_all.pth')