作者：西北小生_
鏈接：http://www.itdecent.cn/p/1cd6333128a1
來源：簡書
著作權(quán)歸作者所有。商業(yè)轉(zhuǎn)載請聯(lián)系作者獲得授權(quán)，非商業(yè)轉(zhuǎn)載請注明出處。

pytorch保存模型非常簡單，主要有兩種方法：

1. 只保存參數(shù)；（官方推薦）
2. 保存整個模型 (結(jié)構(gòu)+參數(shù))。
   由于保存整個模型將耗費大量的存儲，故官方推薦只保存參數(shù)，然后在建好模型的基礎上加載。本文介紹兩種方法，但只就第一種方法進行舉例詳解。

一、只保存參數(shù)

1.保存

一般地，采用一條語句即可保存參數(shù)：

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

其中model指定義的模型實例變量，如 model=vgg16( ), path是保存參數(shù)的路徑，如 path='./model.pth' , path='./model.tar', path='./model.pkl', 保存參數(shù)的文件一定要有后綴擴展名。

特別地，如果還想保存某一次訓練采用的優(yōu)化器、epochs等信息，可將這些信息組合起來構(gòu)成一個字典，然后將字典保存起來：

state = {'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'epoch': epoch}
torch.save(state, path)

2.加載

針對上述第一種情況，也只需要一句即可加載模型：

model.load_state_dict(torch.load(path))

針對上述第二種以字典形式保存的方法，加載方式如下：

checkpoint = torch.load(path)
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
epoch = checkpoint(['epoch'])

需要注意的是，只保存參數(shù)的方法在加載的時候要事先定義好跟原模型一致的模型，并在該模型的實例對象(假設名為model)上進行加載，即在使用上述加載語句前已經(jīng)有定義了一個和原模型一樣的Net, 并且進行了實例化 model=Net( ) 。

另外，如果每一個epoch或每n個epoch都要保存一次參數(shù)，可設置不同的path，如 path='./model' + str(epoch) +'.pth'，這樣，不同epoch的參數(shù)就能保存在不同的文件中，選擇保存識別率最大的模型參數(shù)也一樣，只需在保存模型語句前加個if判斷語句即可。

下面給出一個具體的例子程序，該程序只保存最新的參數(shù)：

#-*- coding:utf-8 -*-

'''本文件用于舉例說明pytorch保存和加載文件的方法'''

__author__ = 'puxitong from UESTC'

import torch as torch
import torchvision as tv
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torchvision.transforms as transforms
from torchvision.transforms import ToPILImage
import torch.backends.cudnn as cudnn
import datetime
import argparse

# 參數(shù)聲明
batch_size = 32
epochs = 10
WORKERS = 0   # dataloder線程數(shù)
test_flag = True  #測試標志，True時加載保存好的模型進行測試 
ROOT = '/home/pxt/pytorch/cifar'  # MNIST數(shù)據(jù)集保存路徑
log_dir = '/home/pxt/pytorch/logs/cifar_model.pth'  # 模型保存路徑

# 加載MNIST數(shù)據(jù)集
transform = tv.transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])

train_data = tv.datasets.CIFAR10(root=ROOT, train=True, download=True, transform=transform)
test_data = tv.datasets.CIFAR10(root=ROOT, train=False, download=False, transform=transform)

train_load = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=WORKERS)
test_load = torch.utils.data.DataLoader(test_data, batch_size=batch_size, shuffle=False, num_workers=WORKERS)

# 構(gòu)造模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, 3, padding=1)
        self.conv2 = nn.Conv2d(64, 128, 3, padding=1)
        self.conv3 = nn.Conv2d(128, 256, 3, padding=1)
        self.conv4 = nn.Conv2d(256, 256, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(256 * 8 * 8, 1024)
        self.fc2 = nn.Linear(1024, 256)
        self.fc3 = nn.Linear(256, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.pool(F.relu(self.conv2(x)))
        x = F.relu(self.conv3(x))
        x = self.pool(F.relu(self.conv4(x)))
        x = x.view(-1, x.size()[1] * x.size()[2] * x.size()[3])
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

model = Net().cuda()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 模型訓練
def train(model, train_loader, epoch):
    model.train()
    train_loss = 0
    for i, data in enumerate(train_loader, 0):
        x, y = data
        x = x.cuda()
        y = y.cuda()
        optimizer.zero_grad()
        y_hat = model(x)
        loss = criterion(y_hat, y)
        loss.backward()
        optimizer.step()
        train_loss += loss
    loss_mean = train_loss / (i+1)
    print('Train Epoch: {}\t Loss: {:.6f}'.format(epoch, loss_mean.item()))

# 模型測試
def test(model, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for i, data in enumerate(test_loader, 0):
            x, y = data
            x = x.cuda()
            y = y.cuda()
            optimizer.zero_grad()
            y_hat = model(x)
            test_loss += criterion(y_hat, y).item()
            pred = y_hat.max(1, keepdim=True)[1]
            correct += pred.eq(y.view_as(pred)).sum().item()
        test_loss /= (i+1)
        print('Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
            test_loss, correct, len(test_data), 100. * correct / len(test_data)))

def main():

    # 如果test_flag=True,則加載已保存的模型
    if test_flag:
        # 加載保存的模型直接進行測試機驗證，不進行此模塊以后的步驟
        checkpoint = torch.load(log_dir)
        model.load_state_dict(checkpoint['model'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        epochs = checkpoint['epoch']
        test(model, test_load)
        return

    for epoch in range(0, epochs):
        train(model, train_load, epoch)
        test(model, test_load)
        # 保存模型
        state = {'model':model.state_dict(), 'optimizer':optimizer.state_dict(), 'epoch':epoch}
        torch.save(state, log_dir)

if __name__ == '__main__':
    main()

3.在加載的模型基礎上繼續(xù)訓練

在訓練模型的時候可能會因為一些問題導致程序中斷，或者常常需要觀察訓練情況的變化來更改學習率等參數(shù)，這時候就需要加載中斷前保存的模型，并在此基礎上繼續(xù)訓練，這時候只需要對上例中的 main() 函數(shù)做相應的修改即可，修改后的 main() 函數(shù)如下：

def main():

    # 如果test_flag=True,則加載已保存的模型
    if test_flag:
        # 加載保存的模型直接進行測試機驗證，不進行此模塊以后的步驟
        checkpoint = torch.load(log_dir)
        model.load_state_dict(checkpoint['model'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        start_epoch = checkpoint['epoch']
        test(model, test_load)
        return

    # 如果有保存的模型，則加載模型，并在其基礎上繼續(xù)訓練
    if os.path.exists(log_dir):
        checkpoint = torch.load(log_dir)
        model.load_state_dict(checkpoint['model'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        start_epoch = checkpoint['epoch']
        print('加載 epoch {} 成功！'.format(start_epoch))
    else:
        start_epoch = 0
        print('無保存模型，將從頭開始訓練！')

    for epoch in range(start_epoch+1, epochs):
        train(model, train_load, epoch)
        test(model, test_load)
        # 保存模型
        state = {'model':model.state_dict(), 'optimizer':optimizer.state_dict(), 'epoch':epoch}
        torch.save(state, log_dir)

以上方法，如果想在命令行進行操作執(zhí)行，都只需加入argpase模塊參數(shù)即可，相關(guān)方法可參考我的博客

二、保存整個模型

1.保存

torch.save(model, path)

2.加載

model = torch.load(path)

用法可參照上例。