当前位置:   article > 正文

torch笔记十二 | 读取和保存_torch 读取pkl

torch 读取pkl

一、方法

·通过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
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
  • 19
  • 20
  • 21
  • 22
  • 23

三、读写模型

# 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
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
  • 19
  • 20
  • 21
  • 22
  • 23
  • 24
  • 25
  • 26
  • 27
  • 28
  • 29
  • 30
  • 31
  • 32
  • 33

四、实例

1.理解

保存:

保存训练好的整个网络:torch.save(net1,'net1.pth')
只保存网络的参数:torch.save(net1.state_dict(),'net1_params.pth')
  • 1
  • 2

提取:

调用整个训练好的网络:net2 = torch.load('net1.pth')
只调用网络的参数:net3.load_state_dict(torch.load('net1_params.pth'))
  • 1
  • 2

其中,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()
  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14
  • 15
  • 16
  • 17
  • 18
  • 19
  • 20
  • 21
  • 22
  • 23
  • 24
  • 25
  • 26
  • 27
  • 28
  • 29
  • 30
  • 31
  • 32
  • 33
  • 34
  • 35
  • 36
  • 37
  • 38
  • 39
  • 40
  • 41
  • 42
  • 43
  • 44
  • 45
  • 46
  • 47
  • 48
  • 49
  • 50
  • 51
  • 52
  • 53
  • 54
  • 55
  • 56
  • 57
  • 58
  • 59
  • 60
  • 61
  • 62
  • 63
  • 64
  • 65
  • 66
  • 67
  • 68
  • 69
  • 70
  • 71
  • 72
  • 73
声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/从前慢现在也慢/article/detail/772749
推荐阅读
相关标签
  

闽ICP备14008679号