pytorch模型的保存加載與續(xù)訓(xùn)練詳解

更新時(shí)間：2022年11月10日 09:54:05 作者：禿頭小蘇

這篇文章主要為大家介紹了pytorch模型的保存加載與續(xù)訓(xùn)練詳解，有需要的朋友可以借鑒參考下，希望能夠有所幫助，祝大家多多進(jìn)步，早日升職加薪

前面

最近，看到不少小伙伴問(wèn)pytorch如何保存和加載模型，其實(shí)這部分pytorch官網(wǎng)介紹的也是很清楚的，感興趣的點(diǎn)擊了解詳情??????

但是肯定有很多人是不愿意看官網(wǎng)的，所以我還是花一篇文章來(lái)為大家介紹介紹。當(dāng)然了，在介紹中我會(huì)加入自己的一些理解，讓大家有一個(gè)更深的認(rèn)識(shí)。如果準(zhǔn)備好了的話(huà)，就讓我們開(kāi)始吧。???

模型保存與加載

pytorch中介紹了幾種不同的模型保存和加載方式，我會(huì)在下文一一為大家介紹。首先先讓我們來(lái)隨便定義一個(gè)模型，如下：【用的是pytorch官網(wǎng)的例子】

# 模型定義
class TheModelClass(nn.Module):
    def __init__(self):
        super(TheModelClass, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)
    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

定義好模型結(jié)構(gòu)后，我們可以實(shí)例化這個(gè)模型：

#模型初始化
model = TheModelClass()

模型初始化過(guò)后，我們就一起來(lái)看看模型保存和加載的方式吧。??????

方式1

方式1是官方推薦的一種方式，我們直接來(lái)看代碼好了，如下：

# 保存模型
torch.save(model.state_dict(), './model/model_state_dict.pth')

該方法后面的參數(shù)'./model/model_state_dict.pth'為模型的保存路徑，模型后綴名官方推薦使用.pth和.pt，當(dāng)然了，你取別的后綴名也是完全可行的。???

介紹了模型的保存，下面就來(lái)看看方式1是如何加載模型的?！具@里我說(shuō)明一點(diǎn)，模型保存往往是在訓(xùn)練中進(jìn)行的，而模型加載多數(shù)用在模型推理中，它們存在兩個(gè)文件中，故我們?cè)谕评磉^(guò)程中要先實(shí)列化模型】

# 加載模型
model_test1 = TheModelClass()   # 加載模型時(shí)應(yīng)先實(shí)例化模型
# load_state_dict()函數(shù)接收一個(gè)字典，所以不能直接將'./model/model_state_dict.pth'傳入，而是先使用load函數(shù)將保存的模型參數(shù)反序列化
model_test1.load_state_dict(torch.load('./model/model_state_dict.pth'))
model_test1.eval()    # 模型推理時(shí)設(shè)置

在上述的代碼注釋中我有寫(xiě)到，我們使用load_state_dict()加載模型時(shí)先需要使用load方法將保存的模型參數(shù)==反序列化==，load后的結(jié)果是一個(gè)字典，這時(shí)就可以通過(guò)load_state_dict()方法來(lái)加載了。

這里我來(lái)簡(jiǎn)單說(shuō)一下我理解的反序列化，其和序列化是相對(duì)應(yīng)的一個(gè)概念。序列化就是把內(nèi)存中的數(shù)據(jù)保存到磁盤(pán)中，像我們使用torch.save()方法保存模型就是序列化；而反序列化則是將硬盤(pán)中的數(shù)據(jù)加載到內(nèi)存當(dāng)中，顯然我們加載模型的過(guò)程就是反序列化過(guò)程?！敬笾碌囊馑既缦聢D所示，偶然在水群的時(shí)候看到一個(gè)畫(huà)圖軟件，是不是還挺好看的??????】

方式2

方式2非常簡(jiǎn)單，直接上代碼：

# 保存模型
torch.save(model, './model/model.pt')    #這里我們保存模型的后綴名取.pt

# 加載模型
model_test2 = torch.load('./model/model.pt')     
model_test2.eval()   # 模型推理時(shí)設(shè)置

但是這種方式是不推薦使用的，因?yàn)槟闶褂眠@種方式保存模型，然后再加載時(shí)會(huì)遇到各種各樣的錯(cuò)誤。為了加深大家理解，我們來(lái)看這樣的一個(gè)例子。文件的結(jié)構(gòu)如下圖所示：

models.py文件中存儲(chǔ)的是模型的定義，其位于文件夾models下。save_model.py文件中寫(xiě)的是保存模型的代碼，如下：

from models.models import TheModelClass
from torch import optim
import torch
#模型初始化
model = TheModelClass()
# 初始化優(yōu)化器
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
# ## 保存加載方式2——save/load
# # 保存模型
# torch.save(models, './models/models.pt')

執(zhí)行此文件后，會(huì)生成models.pt文件，我們?cè)趫?zhí)行load_mode.py文件即可實(shí)現(xiàn)加載，load_mode.py內(nèi)容如下：

from models.models import TheModelClass
import torch
## 加載方式2
# 加載模型
model_test2 = TheModelClass()
model_test2 = torch.load('./models/models.pt')     
model_test2.eval()   # 模型推理時(shí)設(shè)置
print(model_test2)

此時(shí)我們可以正常加載。但如果我們將models文件夾修改為model，如下：

此時(shí)我們?cè)谑褂萌缦麓a加載模型的話(huà)就會(huì)出現(xiàn)錯(cuò)誤：

from models.models import TheModelClass
import torch
## 加載方式2
# 加載模型
model_test2 = TheModelClass()
model_test2 = torch.load('./model/models.pt')     #這里需要修改一下文件路徑  
model_test2.eval()   # 模型推理時(shí)設(shè)置
print(model_test2)

出現(xiàn)這種錯(cuò)誤的原因是使用方式2進(jìn)行模型保存的時(shí)候會(huì)把模型結(jié)構(gòu)定義文件路徑記錄下來(lái)，加載的時(shí)候就會(huì)根據(jù)路徑解析它然后裝載參數(shù)；當(dāng)把模型定義文件路徑修改以后，使用torch.load(path)就會(huì)報(bào)錯(cuò)。

其實(shí)使用方式2進(jìn)行模型的保存和加載還會(huì)存在各種問(wèn)題，感興趣的可以看看這篇博文。總之，在我們今后的使用中，盡量不要用方式2來(lái)加載模型。??????

方式3

pytorch還為我們提供了一種模型保存與加載的方式——checkpoint。這種方式保存的是一個(gè)字典，如果我們程序在運(yùn)行中由于某種原因異常中止，那么這種方式可以很方便的讓我們接著上次訓(xùn)練，正因?yàn)檫@樣，我非常推薦大家使用這種方式進(jìn)行模型的保存與加載。下面就讓我們一起來(lái)看看方式3是如何使用的吧！??！??????

首先，我們同樣使用torch.save來(lái)保存模型，但是這里保存的是一個(gè)字典，里面可以填入你需要保存的參數(shù)，如下：

# 保存checkpoint
torch.save({
            'epoch':epoch,
            'model_state_dict':model.state_dict(),
            'optimizer_state_dict':optimizer.state_dict(),
            'loss':loss
            }, './model/model_checkpoint.tar'    #這里的后綴名官方推薦使用.tar
            )

接著我們來(lái)看看如何加載checkpoint，代碼如下：

# 加載checkpoint
model_checkpoint = TheModelClass()
optimizer =  optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
checkpoint = torch.load('./model/model_checkpoint.tar')    # 先反序列化模型
model_checkpoint.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

看了我上文的介紹，大家是否知道如何使用checkpoint了呢，我想大家都會(huì)覺(jué)得這個(gè)不是很難，但要自己寫(xiě)可能還是不好把握，那么第一次就讓我來(lái)帶領(lǐng)大家看看如何在代碼中使用checkpoint吧?。?！??????

這節(jié)我采用cifar10數(shù)據(jù)集實(shí)現(xiàn)物體分類(lèi)的例子，我的這篇博文對(duì)其進(jìn)行了詳細(xì)介紹，那么這里介紹checkpoint我將利用這個(gè)demo來(lái)為大家講解。首先我們直接來(lái)看模型保存的完整代碼，如下：

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import torchvision
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#1、準(zhǔn)備數(shù)據(jù)集
train_dataset = torchvision.datasets.CIFAR10("./data", train=True, transform=torchvision.transforms.ToTensor(), download= True)
test_dataset = torchvision.datasets.CIFAR10("./data", train=False, transform=torchvision.transforms.ToTensor(), download= True)
#2、加載數(shù)據(jù)集
train_dataset_loader = DataLoader(dataset=train_dataset, batch_size=100)
test_dataset_loader = DataLoader(dataset=test_dataset, batch_size=100)
#3、搭建神經(jīng)網(wǎng)絡(luò)
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.model1 = nn.Sequential(
            nn.Conv2d(3, 32, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(1024, 64),
            nn.Linear(64, 10)
        )
    def forward(self, input):
        input = self.model1(input)
        return input
#4、創(chuàng)建網(wǎng)絡(luò)模型
net = Net()
#5、設(shè)置損失函數(shù)、優(yōu)化器
#損失函數(shù)
loss_fun = nn.CrossEntropyLoss()   #交叉熵
loss_fun = loss_fun.to(device)
#優(yōu)化器
learning_rate = 1e-2
optimizer = torch.optim.SGD(net.parameters(), learning_rate)   #SGD:梯度下降算法
#6、設(shè)置網(wǎng)絡(luò)訓(xùn)練中的一些參數(shù)
total_train_step = 0   #記錄總計(jì)訓(xùn)練次數(shù)
total_test_step = 0    #記錄總計(jì)測(cè)試次數(shù)
Max_epoch = 10    #設(shè)計(jì)訓(xùn)練輪數(shù)
#7、開(kāi)始進(jìn)行訓(xùn)練
for epoch in range(Max_epoch):
    print("---第{}輪訓(xùn)練開(kāi)始---".format(epoch))
    net.train()     #開(kāi)始訓(xùn)練，不是必須的，在網(wǎng)絡(luò)中有BN，dropout時(shí)需要
    #由于訓(xùn)練集數(shù)據(jù)較多，這里我沒(méi)用訓(xùn)練集訓(xùn)練，而是采用測(cè)試集（test_dataset_loader）當(dāng)訓(xùn)練集，但思想是一致的
    for data in test_dataset_loader:      
        imgs, targets = data
        targets = targets.to(device)
        outputs = net(imgs)
        #比較輸出與真實(shí)值，計(jì)算Loss
        loss = loss_fun(outputs, targets)
        #反向傳播，調(diào)整參數(shù)
        optimizer.zero_grad()    #每次讓梯度重置
        loss.backward()
        optimizer.step()
        total_train_step += 1
        if total_train_step % 50 == 0:
            print("---第{}次訓(xùn)練結(jié)束, Loss:{})".format(total_train_step, loss.item()))
    if (epoch+1) % 2 == 0:
        # 保存checkpoint
        torch.save({
            'epoch': epoch,
            'model_state_dict': net.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss
        }, './model/model_checkpoint_epoch_{}.tar'.format(epoch)  # 這里的后綴名官方推薦使用.tar
        )
    if epoch > 5:
        print("---意外中斷---")
        break

整個(gè)流程和這篇文章基本一致，不清楚的建議先花幾分鐘閱讀一下哈。??????主要區(qū)別就是在最后保存模型的時(shí)候我使用了checkpoint進(jìn)行保存，且兩個(gè)epoch保存一次。當(dāng)epoch=6時(shí)，我設(shè)置了一個(gè)break模擬程序意外中斷，中斷后可以來(lái)看一下終端的輸出信息，如下圖所示：

我們可以看到在進(jìn)行第6輪循環(huán)時(shí)，程序中斷了，此時(shí)最新的保存的模型是第五次訓(xùn)練結(jié)果，如下：

同時(shí)注意到第5次訓(xùn)練結(jié)束的loss在2.0左右，如果我們下次接著訓(xùn)練，損失應(yīng)該是在2.0附近。??????

好了，上面由于一些糟糕的原因?qū)е鲁绦蛑袛嗔?，現(xiàn)在我想接著上次訓(xùn)練的結(jié)果繼續(xù)訓(xùn)練，我該怎么辦呢？代碼如下：

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import torchvision
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#1、準(zhǔn)備數(shù)據(jù)集
train_dataset = torchvision.datasets.CIFAR10("./data", train=True, transform=torchvision.transforms.ToTensor(), download= True)
test_dataset = torchvision.datasets.CIFAR10("./data", train=False, transform=torchvision.transforms.ToTensor(), download= True)
#2、加載數(shù)據(jù)集
train_dataset_loader = DataLoader(dataset=train_dataset, batch_size=100)
test_dataset_loader = DataLoader(dataset=test_dataset, batch_size=100)
#3、搭建神經(jīng)網(wǎng)絡(luò)
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.model1 = nn.Sequential(
            nn.Conv2d(3, 32, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(1024, 64),
            nn.Linear(64, 10)
        )
    def forward(self, input):
        input = self.model1(input)
        return input
#4、創(chuàng)建網(wǎng)絡(luò)模型
net = Net()
#5、設(shè)置損失函數(shù)、優(yōu)化器
#損失函數(shù)
loss_fun = nn.CrossEntropyLoss()   #交叉熵
loss_fun = loss_fun.to(device)
#優(yōu)化器
learning_rate = 1e-2
optimizer = torch.optim.SGD(net.parameters(), learning_rate)   #SGD:梯度下降算法
#6、設(shè)置網(wǎng)絡(luò)訓(xùn)練中的一些參數(shù)
total_train_step = 0   #記錄總計(jì)訓(xùn)練次數(shù)
total_test_step = 0    #記錄總計(jì)測(cè)試次數(shù)
Max_epoch = 10    #設(shè)計(jì)訓(xùn)練輪數(shù)
##########################################################################################
# 加載checkpoint
checkpoint = torch.load('./model/model_checkpoint_epoch_5.tar')    # 先反序列化模型
net.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
start_epoch = checkpoint['epoch']
loss = checkpoint['loss']
##########################################################################################
#7、開(kāi)始進(jìn)行訓(xùn)練
for epoch in range(start_epoch+1, Max_epoch):
    print("---第{}輪訓(xùn)練開(kāi)始---".format(epoch))
    net.train()     #開(kāi)始訓(xùn)練，不是必須的，在網(wǎng)絡(luò)中有BN，dropout時(shí)需要
    for data in test_dataset_loader:
        imgs, targets = data
        targets = targets.to(device)
        outputs = net(imgs)
        #比較輸出與真實(shí)值，計(jì)算Loss
        loss = loss_fun(outputs, targets)
        #反向傳播，調(diào)整參數(shù)
        optimizer.zero_grad()    #每次讓梯度重置
        loss.backward()
        optimizer.step()
        total_train_step += 1
        if total_train_step % 50 == 0:
            print("---第{}次訓(xùn)練結(jié)束, Loss:{})".format(total_train_step, loss.item()))
    if (epoch+1) % 2 == 0:
        # 保存checkpoint
        torch.save({
            'epoch': epoch,
            'model_state_dict': net.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss
        }, './model/model_checkpoint_epoch_{}.tar'.format(epoch)  # 這里的后綴名官方推薦使用.tar
        )

這里的代碼相較之前的多了一個(gè)加載checkpoint的過(guò)程，我將其截取出來(lái)，如下圖所示：

通過(guò)加載checkpoint我們就保存了之前訓(xùn)練的參數(shù)，進(jìn)而實(shí)現(xiàn)斷點(diǎn)續(xù)訓(xùn)練，我們直接來(lái)看執(zhí)行此代碼的結(jié)果，如下圖所示：

從上圖可以看出我們的訓(xùn)練是從第6輪開(kāi)始的，并且初始的loss為1.99，和2.0接近。這就說(shuō)明了我們已經(jīng)實(shí)現(xiàn)了中斷后恢復(fù)訓(xùn)練的操作。

????????????????????????????????????????

這里我簡(jiǎn)單的說(shuō)兩句，上文介紹checkpoint的用法時(shí)，訓(xùn)練中斷和訓(xùn)練恢復(fù)我是放在兩個(gè)文件中的進(jìn)行的，但是在實(shí)際中我們肯定是在一個(gè)文件中運(yùn)行，那這該怎么辦呢？其實(shí)方法很簡(jiǎn)單啦，我們只需要設(shè)置一個(gè)if條件將加載checkpoint的部分放在訓(xùn)練文件中，然后設(shè)置一個(gè)參數(shù)來(lái)控制if條件的執(zhí)行即可。具體細(xì)節(jié)我就不給大家介紹了，如果有不明白的評(píng)論區(qū)見(jiàn)吧！?。????????

????????????????????????????????????????