在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)