PyTorch：保存函数pytorch.save()和加载函数pytorch.load()的使用

在Pytorch中，保存和加载操作通常是成对发生的，我们使用pytorch.save()和pytorch.load()函数来分别完成这两个功能。在训练模型时，我们可以将训练得到的模型、模型的权重、优化器的参数甚至一些Tensor都保存下来，在以后使用到该模型时直接加载，而不需要每次使用时都重新训练一遍这些参数。

一、模型的保存和加载

首先创建一个简单的模型。

import torch
from torch import nn
from collections import OrderedDict

net = nn.Sequential(OrderedDict([
    ('linear1', nn.Linear(2, 10)),
    ('relu', nn.ReLU()),
    ('linear2', nn.Linear(10, 1)),
    ('sigmoid', nn.Sigmoid())
]))

# 省略模型训练过程...

我们可以将模型结构和训练得到的参数一起保存下来。

torch.save(net, 'net.pt')

当我们使用时，直接从net.pt文件中加载这个模型即可。

net_l = torch.load('net.pt')

我们可以将样本输入到这个加载得到的模型中并获得输出。

output = net_l(input)