pytorch保存模型的方式有两种

     ①将整个网络都都保存下来

           保存整个神经网络的的结构信息和模型参数信息,save的对象是网络net

           后缀一般命名为.pkl

     ②仅保存和加载模型参数(推荐使用这样的方法)

           当需要为预测保存一个模型的时候,只需要保存训练模型的可学习参数即可。只保存神经网络的训练模型参数,save的对象是net.state_dict()

           后缀一般命名为 .pt 或者 .pth

但是记住在进行预测之前,必须调用 model.eval() 方法来将 dropout 和 batch normalization 层设置为验证模型。否则,只会生成前后不一致的预测结果。
  

加载方式

①加载模型时通过torch.load('.pth')直接初始化新的神经网络对象

②需要首先导入对应的网络,再通过net.load_state_dict(torch.load('.pth'))完成模型参数的加载

net = Net() # 保存和加载整个模型 torch.save(net, 'model.pkl') model = torch.load('model.pkl')# 仅保存和加载模型参数(推荐使用) torch.save(net.state_dict(), 'params.pkl') net.load_state_dict(torch.load('params.pkl'))

(后缀是pth还是pkl还是ckpt.mdl都无所谓的,你命名是啥就是啥)

记住这两种方式都是需要为预测保存模型的时候,即你已经把模型train完了,需要保存下来以备预测使用

但是如果在train的过程中保存,即你训练还没完不得已先结束或者断电断网了之类的,那就需要加载和保存一个通用的检查点(Checkpoint),也就是断点续传

直接读取,而不加载到模型中

import torch
checkpoint = torch.load('output/trained_model.pth')
print(checkpoint)
然后如果删掉几个key值元素再保存的话
del checkpoint['conv']
torch.save(checkpoint, 'work_dirs/self_annotated_IN_backbone_last_2_stages/epoch_12_without_char-bbox-head-fc-cls.pth')


甚至从其他模型中取一部分,写入到当前模型中再让程序加载也是ok的

import torch
checkpoint = torch.load('work_dirs/self_annotated_IN_backbone_last_2_stages/epoch_12_without_char-bbox-head-fc-cls.pth')
checkpoint_landmark = torch.load('work_dirs/self_annotated_IN_backbone_last_2_stages_landmark_3-layer_conv_10w+3w-data_divloss*50_6-epoch_feed-testdata-3k/epoch_6_only_with_landmark-conv.pth')
checkpoint['state_dict'].update(checkpoint_landmark['state_dict'])

torch.save(checkpoint, 'work_dirs/self_annotated_IN_backbone_last_2_stages/epoch_12_without_char-bbox-head-fc-cls_with_landmark.pth')

注意加载权重是要key名称匹配的

即使是相同结构,名称不同也是没法load的

删除部分模型

其实删除模型层就是替换模型部分层,我们只需要写一个空层替换局可以了,
如下:

#删除最后一个层,或者说删除任意一个层 # 写一个空的层 classifier = nn.Sequential() #然后替换一下就可以了 CNN_shuffmodel.fc = classifier

删除层就是这么简单
 

模型的训练过程中加载和保存Checkpoint(断点续传)

      当保存一个通用的检查点(checkpoint)时,无论是用于继续训练还是预测,都需要保存更多的信息,不仅仅是 state_dict ,比如说优化器的 state_dict 也是非常重要的,它包含了用于模型训练时需要更新的参数和缓存信息,还可以保存的信息包括 epoch,即中断训练的批次,最后一次的训练 loss,额外的 torch.nn.Embedding 层等等。

上述保存代码就是介绍了如何保存这么多种信息,通过用一个字典来进行组织,然后继续调用 torch.save 方法

保存

checkpoint = {
        "net": model.state_dict(),
        'optimizer':optimizer.state_dict(),
        "epoch": epoch
    }

torch.save(checkpoint, 'checkpoints/ckpt_%s.pth' %(str(epoch)))
加载
checkpoint = torch.load('checkpoints/ckpt_100.pth')  # 加载断点

model.load_state_dict(checkpoint['net'])  # 加载模型可学习参数
optimizer.load_state_dict(checkpoint['optimizer'])  # 加载优化器参数
start_epoch = checkpoint['epoch']  # 设置开始的epoch
#加载之后再训练就直接从start_epoch+1开始了
更详细的见

加载完后,根据后续步骤,调用 model.eval() 用于预测,model.train() 用于恢复训练。

在GPU上保存模型,在 CPU 上加载模型

保存模型的示例代码

torch.save(model.state_dict(), PATH)

加载模型的示例代码

device = torch.device('cpu') model = TheModelClass(*args, **kwargs) model.load_state_dict(torch.load(PATH, map_location=device))

在 CPU 上加载在 GPU 上训练的模型,必须在调用 torch.load() 的时候,设置参数 map_location ,指定采用的设备是 torch.device('cpu'),这个做法会将张量都重新映射到 CPU 上。

在GPU上保存模型,在 GPU 上加载模型

保存模型的示例代码

torch.save(model.state_dict(), PATH)

加载模型的示例代码

device = torch.device('cuda') model = TheModelClass(*args, **kwargs) model.load_state_dict(torch.load(PATH) model.to(device) # Make sure to call input = input.to(device) on any input tensors that you feed to the model

在 GPU 上训练和加载模型,调用 torch.load() 加载模型后,还需要采用 model.to(torch.device('cuda')),将模型调用到 GPU 上,并且后续输入的张量都需要确保是在 GPU 上使用的,即也需要采用 my_tensor.to(device)

在CPU上保存,在GPU上加载模型

保存模型的示例代码

torch.save(model.state_dict(), PATH)

加载模型的示例代码

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location="cuda:0"))  # Choose whatever GPU device number you want
model.to(device)
# Make sure to call input = input.to(device) on any input tensors that you feed to the model

这次是 CPU 上训练模型,但在 GPU 上加载模型使用,那么就需要通过参数 map_location 指定设备。然后继续记得调用 model.to(torch.device('cuda'))

序列化 & 反序列化

PyTorch中模型,参数的保存

序列化——从内存到硬盘

反序列化——加载,从硬盘到内存

实操

我们使用vgg19模型 

vgg19.py
import os
import torch
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision import transforms
from torch import nn,optim
#from lenet5 import Lenet5
from vgg import VGG19
from vgg import VGG34
 
def main():
    batchsz = 32
 
    cifar_train = datasets.CIFAR10('dataset/', train=True, transform=transforms.Compose([
        transforms.Resize((32, 32)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485,0.456,0.406],
        					 std=[0.229,0.224,0.225])
    ]), download=True)
    cifar_train = DataLoader(cifar_train, batch_size=batchsz, shuffle=True)

    cifar_test = datasets.CIFAR10('dataset/', train=False, transform=transforms.Compose([
        transforms.Resize((32, 32)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485,0.456,0.406],
        					 std=[0.229,0.224,0.225])
    ]), download=True)
    cifar_test = DataLoader(cifar_test, batch_size=batchsz, shuffle=True)
 
    x, label = iter(cifar_train).next()
    print('x:', x.shape, 'label:', label.shape)
 
    device = torch.device('cuda')
    model = VGG19().to(device)
 
    criton = nn.CrossEntropyLoss().to(device)  #包含了softmax
    optimizer = optim.Adam(model.parameters(),lr=1e-3)
    print(model)
    
    if os.path.exists('model.pkl'):
        model.load_state_dict(torch.load('model.pkl'))
        print('model loaded from model.pkl')
 
    for epoch in range(20):
        model.train()
        for batchidx, (x,label) in enumerate(cifar_train):
            #x:[b,3,32,32]
            #label: [b]
            x,label = x.to(device), label.to(device)
            logits=model(x)
            #logits:[b,10]
            #label:b[b]
            #loss: tensor scalar  长度为0的标量
            loss = criton(logits,label)
 
            #反向传播
            optimizer.zero_grad()  #梯度清零
            loss.backward() #得到新的梯度
            optimizer.step() #走一步,把梯度更新到weight里面去了
 
        print('epoch',epoch, ' loss: ',loss.item())
 
        model.eval()
        #test
        total_correct = 0
        total_num = 0
        for x, label in cifar_test:
            # x:[b,3,32,32]
            # label: [b]
            x,label = x.to(device) ,label.to(device)
 
            #logits:[b,10]
            logits = model(x)
            pred = logits.argmax(dim=1)
 
            total_correct += torch.eq(pred,label).float().sum().item()
            total_num += x.size(0)  #即batch_size
        acc = total_correct / total_num
        print('epoch',epoch, ' acc: ',acc)

        torch.save(model.state_dict(),'model.pkl')
        print('model in epoch',epoch,' saved in model.pkl')
 
 
if __name__ == '__main__':
    main()

模型训练结果

pytorch 保存 list pytorch 保存pkl 在线调用_pytorch 保存 list

现在我们再运行一次,看一下直接加载model后的训练结果

pytorch 保存 list pytorch 保存pkl 在线调用_加载_02

直接就在原来acc 67%的模型上继续训练,上来就70%多了

如果只是想验证一下这个模型,不用训练

import os
import torch
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision import transforms
from torch import nn,optim
#from lenet5 import Lenet5
from vgg import VGG19
from vgg import VGG34
 
def main():
    batchsz = 32
 
    cifar_train = datasets.CIFAR10('dataset/', train=True, transform=transforms.Compose([
        transforms.Resize((32, 32)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485,0.456,0.406],
        					 std=[0.229,0.224,0.225])
    ]), download=True)
    cifar_train = DataLoader(cifar_train, batch_size=batchsz, shuffle=True)

    cifar_test = datasets.CIFAR10('dataset/', train=False, transform=transforms.Compose([
        transforms.Resize((32, 32)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485,0.456,0.406],
        					 std=[0.229,0.224,0.225])
    ]), download=True)
    cifar_test = DataLoader(cifar_test, batch_size=batchsz, shuffle=True)
 
    x, label = iter(cifar_train).next()
    print('x:', x.shape, 'label:', label.shape)
 
    device = torch.device('cuda')
    model = VGG19().to(device)
 
    criton = nn.CrossEntropyLoss().to(device)  #包含了softmax
    optimizer = optim.Adam(model.parameters(),lr=1e-3)
    print(model)
    
    if os.path.exists('model.pkl'):
        model.load_state_dict(torch.load('model.pkl'))
        print('model loaded from model.pkl')
 
    
    model.eval()
    #test
    total_correct = 0
    total_num = 0
    for x, label in cifar_test:
        x,label = x.to(device) ,label.to(device)
        logits = model(x)
        pred = logits.argmax(dim=1)
        total_correct += torch.eq(pred,label).float().sum().item()
        total_num += x.size(0)  #即batch_size
    acc = total_correct / total_num
    print('acc: ',acc)
 
 
if __name__ == '__main__':
    main()

pytorch 保存 list pytorch 保存pkl 在线调用_ci_03