一、方法
·通过save函数和load函数可以很方便地读写Tensor。
·通过save函数和load_state_dict函数可以很方便地读写模型的参数。
二、读写tensor
import torch
import torch.nn as nn
# 1.创建x,并将其存入文件名为x.pt的文件里
x = torch.ones(5)
torch.save(x, 'x.pt')
# 2.将数据从存储的文件读回内存
y = torch.load('x.pt')
print(y)
# 3.保存提取列表
x1 = torch.ones(3)
x2 = torch.ones(2)
x_list = [x1, x2]
torch.save(x_list, 'x_list.pt')
y_list = torch.load('x_list.pt')
print(y_list)
# 4.存储并读取一个从字符串映射到Tensor的字典
torch.save({'x1': x1, 'x2': x2}, 'x_dict.pt')
y_dict = torch.load('x_dict.pt')
print(y_dict)
三、读写模型
# 1.只保存和加载模型参数
torch.save(model.state_dict(), PATH) # 推荐的文件后缀名是pt或pth
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
# 2.保存和加载整个模型
torch.save(model, PATH)
model = torch.load(PATH)
# 3. 例子1:
net = torch.nn.Sequential(torch.nn.Linear(3, 1))
X = torch.randn(2, 3)
Y = net(X)
# 保存模型参数
PATH = "./net.pt"
torch.save(net.state_dict(), PATH)
# 加载模型参数
net2 = torch.nn.Sequential(torch.nn.Linear(3, 1)) # 需要先搭建网络结构
net2.load_state_dict(torch.load(PATH))
Y2 = net2(X)
print(Y2 == Y) # True
# 4.例子2:
net = torch.nn.Sequential(torch.nn.Linear(3, 1))
X = torch.randn(2, 3)
Y = net(X)
PATH = "./net1.pt"
torch.save(net, PATH)
net2 = torch.load(PATH) # 不需要重新搭建网络
Y2 = net2(X)
print(Y2 == Y) # True
PS:
四、实例
1.理解
保存:
保存训练好的整个网络:torch.save(net1,'net1.pth')
只保存网络的参数:torch.save(net1.state_dict(),'net1_params.pth')
提取:
调用整个训练好的网络:net2 = torch.load('net1.pth')
只调用网络的参数:net3.load_state_dict(torch.load('net1_params.pth'))
其中,net1是训练好的网络的名称,’ '内部是保存的文件名称(后缀是.pth或者.pkl)
注意:
(1)只保存网络中的参数速度快, 占内存少,推荐使用。但是调用网络的参数时,新网络需要自己定义网络,再使用上面的调用指令。其中的参数名称与结构要与保存模型中的一致,可以是部分网络比如只使用VGG的前几层,相对灵活,便于对网络进行修改。
(2)调用整个网络则无需自定义网络可以直接调用。保存时已把网络结构保存,比较死板,不能调整网络结构。
2.实例
2.1 实验结果
以我的[pytorch学习笔记二]数据的拟合为例,将训练好的net1网络保存好,net2调用的整个网络,net3只调用了net1的参数,最终拟合的效果一模一样。
2.2完整代码
# 1.导入必要的模块
import torch
import torch.nn.functional as F # F中包含很多函数比如激励函数
import matplotlib.pyplot as plt #用于绘图
# 2.生成要拟合的数据点
x = torch.unsqueeze(torch.linspace(-1,1,100),dim=1)
y = x.pow(2) + 0.2 * torch.rand(x.size())
# 3.搭建、训练、保存网络net1
def save():
net1 = torch.nn.Sequential(
torch.nn.Linear(1, 8),
torch.nn.ReLU(),
torch.nn.Linear(8, 1)
)
optimizer = torch.optim.SGD(net1.parameters(), lr=0.5)
loss_func = torch.nn.MSELoss()
for t in range(100):
prediction = net1(x)
loss = loss_func(prediction, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 保存训练好的整个网络
torch.save(net1, 'net1.pth')
# 只保存训练好的网络的参数
torch.save(net1.state_dict(), 'net1_params.pth')
# 绘图,可视化net1的训练结果
plt.figure(1,figsize=(10,3))
plt.subplot(131)
plt.title('net1')
plt.scatter(x.data.numpy(), y.data.numpy()) # 打印原始数据散点图
plt.plot(x.data.numpy(), prediction.data.numpy(), 'r-', lw=5)
# 调用整个网络(包括计算图、参数等等)
def restore_net():
net2 = torch.load('net1.pth')
# 调用网络时,要设置输入
prediction = net2(x)
plt.subplot(132)
plt.title('net2')
plt.scatter(x.data.numpy(), y.data.numpy())
plt.plot(x.data.numpy(), prediction.data.numpy(), 'r-', lw=5)
# 只调用网络的参数
def restore_params():
# 只调用网络的参数时,需要提前搭建和net1网络相同的架构
net3 = torch.nn.Sequential(
torch.nn.Linear(1, 8),
torch.nn.ReLU(),
torch.nn.Linear(8, 1)
)
net3.load_state_dict(torch.load('net1_params.pth'))
prediction = net3(x)
plt.subplot(133)
plt.title('net3')
plt.scatter(x.data.numpy(), y.data.numpy()) # 打印原始数据散点图
plt.plot(x.data.numpy(), prediction.data.numpy(), 'r-', lw=5)
plt.show()
save()
restore_net()
restore_params()