torch笔记十二 | 读取和保存

一、方法

·通过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

四、实例

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