【深度学习DL-PyTorch】六、保存和加载模型

阿里云双11来了!从本博客参与阿里云,服务器最低只要86元/年!

使用 PyTorch 保存和加载模型。

一、 训练网络

model = fc_model.Network(784, 10, [512, 256, 128])
criterion = nn.NLLLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
fc_model.train(model, trainloader, testloader, criterion, optimizer, epochs=2)

二、 保存和加载网络

每次需要使用网络时都去训练它不太现实,也很不方便。我们可以保存训练过的网络,之后加载这些网络来继续训练或用它们进行预测。
PyTorch 网络的参数保存在模型的 state_dict 中。可以看到这个状态字典包含每个层级的权重和偏差矩阵。

print("Our model: \n\n", model, '\n')
print("The state dict keys: \n\n", model.state_dict().keys())

最简单的方法是使用torch.save 保存状态字典。例如,我们可以将其保存到文件 ‘checkpoint.pth’ 中。

torch.save(model.state_dict(), 'checkpoint.pth')

然后,使用 torch.load 加载这个状态字典。

state_dict = torch.load('checkpoint.pth')
print(state_dict.keys())

要将状态字典加载到神经网络中,需要执行 model.load_state_dict(state_dict)

model.load_state_dict(state_dict) 

只有模型结构和检查点的结构完全一样时,状态字典才能加载成功哦。如果我在创建模型时使用了不同的结构,便无法顺利加载。

# Try this
model = fc_model.Network(784, 10, [400, 200, 100])
# This will throw an error because the tensor sizes are wrong!
model.load_state_dict(state_dict)

这就是说,我们需要重新构建和训练时完全一样的模型。我们需要将模型架构信息与状态字典一起保存在检查点里。所以,你需要创建一个字典,其中包含完全重新构建模型所需的所有信息。

checkpoint = {'input_size': 784,
              'output_size': 10,
              'hidden_layers': [each.out_features for each in model.hidden_layers],
              'state_dict': model.state_dict()}
torch.save(checkpoint, 'checkpoint.pth')

现在,检查点中包含了重建训练模型所需的全部信息。你可以随意将它编写为函数。同样,我们可以编写一个函数来加载检查点。

def load_checkpoint(filepath):
    checkpoint = torch.load(filepath)
    model = fc_model.Network(checkpoint['input_size'],
                             checkpoint['output_size'],
                             checkpoint['hidden_layers'])
    model.load_state_dict(checkpoint['state_dict'])
    
    return model
model = load_checkpoint('checkpoint.pth')
print(model)

https://www.jianshu.com/p/7f301fc4c741

Python量化投资网携手4326手游为资深游戏玩家推荐:《《纯白魔女》:【转自投石姬】系统介绍-仇恨机制说明

「点点赞赏,手留余香」

    还没有人赞赏,快来当第一个赞赏的人吧!
0 条回复 A 作者 M 管理员
    所有的伟大,都源于一个勇敢的开始!
欢迎您,新朋友,感谢参与互动!欢迎您 {{author}},您在本站有{{commentsCount}}条评论