pytorch深度学习存储权重的2种方法

(1).npy文件

存储数组，可用于存储深度学习训练得的权值矩阵；该方法适合底层代码,在不用torch框架的情况下,从0编写代码时使用.

如下为单层的线性回归的部分代码:

存储权重:

w = net[1].weight.data
b = net[1].bias.data
np.save('w.npy',w)
np.save('b.npy',b)

读取权重,并导入网络中,即可计算预测结果,其中img_arr为输入的数字图片:

    # 读取权重
    w = np.load('w.npy')
    b = np.load('b.npy')
    # 将权重转换为 PyTorch 的张量格式
    w_tensor = torch.from_numpy(w).float()
    b_tensor = torch.from_numpy(b).float()

    #调用模型
    net = nn.Sequential(nn.Linear(784,10))
    net[0].weight.data = w_tensor
    net[0].bias.data = b_tensor
    out = torch.argmax(net(img_arr))

(2) .state_dict()

该方法适合在torch下使用,能够将多层权重存储在一起,更简便快速.

存储权重:

torch.save(net.state_dict(), 'wb_mnist_linear.params')

读取权重:

net = nn.Sequential(nn.Flatten(),nn.Linear(784,10))
net.load_state_dict(torch.load('wb_mnist_linear.params'))
out = torch.argmax(net(img_arr))