torch笔记十二|读取和保存

一、方法

·通过save函数和load函数可以很方便地读写Tensor。

·通过save函数和load_state_dict函数可以很方便地读写模型的参数。

二、读写tensor

import torch
import torch.nn as nn# 1.创建x&＃xff0c;并将其存入文件名为x.pt的文件里
x &＃61; torch.ones(5)
torch.save(x, &＃39;x.pt&＃39;)# 2.将数据从存储的文件读回内存
y &＃61; torch.load(&＃39;x.pt&＃39;)
print(y)# 3.保存提取列表
x1 &＃61; torch.ones(3)
x2 &＃61; torch.ones(2)
x_list &＃61; [x1, x2]
torch.save(x_list, &＃39;x_list.pt&＃39;)
y_list &＃61; torch.load(&＃39;x_list.pt&＃39;)
print(y_list)# 4.存储并读取一个从字符串映射到Tensor的字典
torch.save({&＃39;x1&＃39;: x1, &＃39;x2&＃39;: x2}, &＃39;x_dict.pt&＃39;)
y_dict &＃61; torch.load(&＃39;x_dict.pt&＃39;)
print(y_dict)


三、读写模型

# 1.只保存和加载模型参数
torch.save(model.state_dict(), PATH) # 推荐的文件后缀名是pt或pth
model &＃61; TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))# 2.保存和加载整个模型
torch.save(model, PATH)
model &＃61; torch.load(PATH)# 3. 例子1&＃xff1a;
net &＃61; torch.nn.Sequential(torch.nn.Linear(3, 1))
X &＃61; torch.randn(2, 3)
Y &＃61; net(X)
# 保存模型参数
PATH &＃61; "./net.pt"
torch.save(net.state_dict(), PATH)
# 加载模型参数
net2 &＃61; torch.nn.Sequential(torch.nn.Linear(3, 1)) # 需要先搭建网络结构
net2.load_state_dict(torch.load(PATH))
Y2 &＃61; net2(X)
print(Y2 &＃61;&＃61; Y) # True# 4.例子2&＃xff1a;
net &＃61; torch.nn.Sequential(torch.nn.Linear(3, 1))
X &＃61; torch.randn(2, 3)
Y &＃61; net(X)PATH &＃61; "./net1.pt"
torch.save(net, PATH)net2 &＃61; torch.load(PATH) # 不需要重新搭建网络
Y2 &＃61; net2(X)
print(Y2 &＃61;&＃61; Y) # True


四、实例

1.理解

保存&＃xff1a;

保存训练好的整个网络&＃xff1a;torch.save(net1,&＃39;net1.pth&＃39;)
只保存网络的参数&＃xff1a;torch.save(net1.state_dict(),&＃39;net1_params.pth&＃39;)


提取&＃xff1a;

调用整个训练好的网络&＃xff1a;net2 &＃61; torch.load(&＃39;net1.pth&＃39;)
只调用网络的参数&＃xff1a;net3.load_state_dict(torch.load(&＃39;net1_params.pth&＃39;))


其中&＃xff0c;net1是训练好的网络的名称&＃xff0c;’ &＃39;内部是保存的文件名称&＃xff08;后缀是.pth或者.pkl&＃xff09;

注意&＃xff1a;
&＃xff08;1&＃xff09;只保存网络中的参数速度快, 占内存少&＃xff0c;推荐使用。但是调用网络的参数时&＃xff0c;新网络需要自己定义网络&＃xff0c;再使用上面的调用指令。其中的参数名称与结构要与保存模型中的一致&＃xff0c;可以是部分网络比如只使用VGG的前几层&＃xff0c;相对灵活&＃xff0c;便于对网络进行修改。
&＃xff08;2&＃xff09;调用整个网络则无需自定义网络可以直接调用。保存时已把网络结构保存&＃xff0c;比较死板&＃xff0c;不能调整网络结构。

2.实例

2.1 实验结果

以我的[pytorch学习笔记二]数据的拟合为例&＃xff0c;将训练好的net1网络保存好&＃xff0c;net2调用的整个网络&＃xff0c;net3只调用了net1的参数&＃xff0c;最终拟合的效果一模一样。

2.2完整代码

# 1.导入必要的模块
import torch
import torch.nn.functional as F # F中包含很多函数比如激励函数
import matplotlib.pyplot as plt #用于绘图# 2.生成要拟合的数据点
x &＃61; torch.unsqueeze(torch.linspace(-1,1,100),dim&＃61;1)
y &＃61; x.pow(2) &＃43; 0.2 * torch.rand(x.size())# 3.搭建、训练、保存网络net1
def save():net1 &＃61; torch.nn.Sequential(torch.nn.Linear(1, 8),torch.nn.ReLU(),torch.nn.Linear(8, 1))optimizer &＃61; torch.optim.SGD(net1.parameters(), lr&＃61;0.5)loss_func &＃61; torch.nn.MSELoss()for t in range(100):prediction &＃61; net1(x)loss &＃61; loss_func(prediction, y)optimizer.zero_grad()loss.backward()optimizer.step()# 保存训练好的整个网络torch.save(net1, &＃39;net1.pth&＃39;)# 只保存训练好的网络的参数torch.save(net1.state_dict(), &＃39;net1_params.pth&＃39;)# 绘图&＃xff0c;可视化net1的训练结果plt.figure(1,figsize&＃61;(10,3))plt.subplot(131)plt.title(&＃39;net1&＃39;)plt.scatter(x.data.numpy(), y.data.numpy()) # 打印原始数据散点图plt.plot(x.data.numpy(), prediction.data.numpy(), &＃39;r-&＃39;, lw&＃61;5)# 调用整个网络&＃xff08;包括计算图、参数等等&＃xff09;
def restore_net():net2 &＃61; torch.load(&＃39;net1.pth&＃39;)# 调用网络时&＃xff0c;要设置输入prediction &＃61; net2(x)plt.subplot(132)plt.title(&＃39;net2&＃39;)plt.scatter(x.data.numpy(), y.data.numpy())plt.plot(x.data.numpy(), prediction.data.numpy(), &＃39;r-&＃39;, lw&＃61;5)# 只调用网络的参数
def restore_params():# 只调用网络的参数时&＃xff0c;需要提前搭建和net1网络相同的架构net3 &＃61; torch.nn.Sequential(torch.nn.Linear(1, 8),torch.nn.ReLU(),torch.nn.Linear(8, 1))net3.load_state_dict(torch.load(&＃39;net1_params.pth&＃39;))prediction &＃61; net3(x)plt.subplot(133)plt.title(&＃39;net3&＃39;)plt.scatter(x.data.numpy(), y.data.numpy()) # 打印原始数据散点图plt.plot(x.data.numpy(), prediction.data.numpy(), &＃39;r-&＃39;, lw&＃61;5)plt.show()save()restore_net()restore_params()