Pytorch_Part7_模型使用

最新推荐文章于 2023-02-14 10:44:51 发布

原创

最新推荐文章于 2023-02-14 10:44:51 发布 · 267 阅读

0 ·

CC 4.0 BY-SA版权

本文介绍了PyTorch中模型的保存与加载，包括序列化、断点续训，详细讲解了模型微调的Transfer Learning和Resnet18迁移学习的应用。此外，还探讨了GPU的使用，包括数据和模型的to操作、GPU设备管理以及多GPU运算的并行机制。

共同贡献PyTorch常见错误与坑汇总文档：《PyTorch常见报错/坑汇总》

一、模型保存与加载

1. 序列化与反序列化

net = LeNet2(classes=2019)

# 法1: 保存整个Module，不仅保存参数，也保存结构
torch.save(net, path)

net_load = torch.load(path_model)	# 网络名称、结构、模型参数、优化器参数均保留

# 法2: 保存模型参数（推荐，占用资源少）
state_dict = net.state_dict()
torch.save(state_dict , path)

net_new = LeNet2(classes=2019)
net_new.load_state_dict(state_dict_load)

2. 断点续训练

保存：

checkpoint = {
   
   
"model_state_dict": net.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"epoch": epoch
}
path_checkpoint = "./checkpoint_{}_epoch.pkl".format(epoch)
torch.save(checkpoint, path_checkpoint)

恢复：

# ============================ step 2/5 模型 ============================
net = LeNet(classes=2)
net.initialize_weights()

# ============================ step 3/5 损失函数 ============================
criterion = nn.CrossEntropyLoss()                                                   # 选择损失函数

# ============================ step 4/5 优化器 ============================
optimizer = optim.SGD(net.parameters(), lr=LR