模型转换(01) : pytorch保存模型

模型转换(01) : pytorch保存模型 1.   模型存储接口

首先我们知道不论是保存模型还是参数都需要用到torch.save()。

对于torch.save()有两种保存方式：

        a. 只保存神经网络的训练模型的参数，save的对象是model.state_dict()；
        b. 既保存整个神经网络的的模型结构又保存模型参数，那么save的对象就是整个模型；

import torch

# 保存模型步骤
torch.save(model, 'net.pth')  # 保存整个神经网络的模型结构以及参数
torch.save(model, 'net.pkl')  # 同上
torch.save(model.state_dict(), 'net_params.pth')  # 只保存模型参数
torch.save(model.state_dict(), 'net_params.pkl')  # 同上

# 加载模型步骤
model = torch.load('net.pth')  # 加载整个神经网络的模型结构以及参数
model = torch.load('net.pkl')  # 同上
model.load_state_dict(torch.load('net_params.pth')) # 仅加载参数
model.load_state_dict(torch.load('net_params.pkl')) # 同上

2.   模型保存加载实例

#  ------------------------  保存模型  --------------------------
torch.save
(
    {
        'epoch': epoch + 1,
        'state_dict': model.state_dict(),
        'best_acc1': best_acc1,
        'optimizer': optimizer.state_dict(),
    }, 
    is_best
)


#  ------------------------  加载模型  --------------------------
# step 1: 模型加载到gpu还是cpu
if args.gpu is None:
    checkpoint = torch.load(args.resume)
else:
    checkpoint = torch.load(args.resume, map_location=args.gpu)
# step 2: 把参数加载到网络
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])

# step 2: 加载模型可能会出现名字对应不上的问题，见下面这种写法：
if args.gpu is not None:
    model.load_state_dict(state_dict)
else:
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        name = k[7:]
        new_state_dict[name] = v
    model.load_state_dict(new_state_dict)

3.   state_dict中内容及层的命名

        a. 对于__init__中使用self定义的变量会使用这个变量名作为存储时的名字。

        self.conv1 = nn.Conv2d(3, 12, kernel_size=3, stride=1, padding=1) 
        卷积层有两个参数：权重和偏移项，上例对应的名称为conv1.weight、conv1.bias

        self.bn1 = nn.BatchNorm2d(12)

        BN层有5个参数：bn1.weight、bn1.bias、bn1.running_mean、bn1.running_var、bn1.num_batches_tracked

        b. 当使用nn.Sequential时会根据传入的list的顺序对其进行编号，从0开始。

        conv1 = nn.Conv2d(3, 12, kernel_size=3, stride=1, padding=1)
        bn1 = nn.BatchNorm2d(12)
        s1 = [conv1, bn1]
        self.stage1 = nn.Sequential(*s1)
        注意此时的conv1和bn1都没有self，stage1有self，由于Sequential将conv1和bn1进行顺序封装，因此conv1会被编号为stage1.0，bn1会被编号为stage1.1，具体结果如下：

        stage1.0.weight、stage1.0.bias

        stage1.1.weight、stage1.1.bias、stage1.1.running_mean、stage1.1.running_var、stage1.1.num_batches_tracked

例: 

content = torch.load("model_best.pth.tar")
content.keys()
# dict_keys(['epoch', 'state_dict', 'best_acc1', 'optimizer'])

content['state_dict']
"""
# []中全是张量
OrderedDict([('conv.0.weight', tensor([], device='cuda:1')),
             ('conv.0.bias', tensor([], device='cuda:1')),
             ('conv.1.weight', tensor([], device='cuda:1')),
             ('conv.1.bias', tensor([ ],device='cuda:1')),
             ('conv.1.running_mean', tensor([],device='cuda:1')),
             ('conv.1.running_var', tensor([], device='cuda:1')),
             ('conv.1.num_batches_tracked', tensor(9675, device='cuda:1')),
             
             # 注意下面这些是在一个block里，有的没加self，在sequential里
             ('block1_1.conv1.0.weight', tensor([], device='cuda:1')),
             ('block1_1.conv1.0.bias', tensor([], device='cuda:1')),
             ('block1_1.conv1.1.weight', tensor([], device='cuda:1')),
             ('block1_1.conv1.1.bias', tensor([],device='cuda:1')),
             ('block1_1.conv1.1.running_mean', tensor([],device='cuda:1')),
             ('block1_1.conv1.1.running_var', tensor([], device='cuda:1')),
             ('block1_1.conv1.1.num_batches_tracked', tensor(9675, device='cuda:1')),
             ('block1_1.conv2.0.weight', tensor([], device='cuda:1')),
             ('block1_1.conv2.0.bias', tensor([],device='cuda:1')),
"""

PyTorch保存网络结构以及参数【 torch.save()、torch.load() 】_Raywit的博客-CSDN博客_torch.save()

pytorch中存储各层权重参数时的命名规则，为什么有些层的名字中带module._南国那片枫叶的博客-CSDN博客_dist.get_rank()