colab pytorch保存模型

Posted xiximayou

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了colab pytorch保存模型相关的知识,希望对你有一定的参考价值。

保存模型总体来说有两种:

第一种:保存训练的模型,之后我们可以继续训练

(1)保存模型

state = { model: model.state_dict(), optimizer:optimizer.state_dict(), epoch: epoch }   
torch.save(state, path)

model.state_dict():模型参数

optimizer.state_dict():优化器

epoch:保存epoch,为了可以接着训练

(2)恢复模型

checkpoint = torch.load(path)
model.load_state_dict(checkpoint[model])
optimizer.load_state_dict(checkpoint[optimizer])
start_epoch = checkpoint[epoch]+1

第二种:保存测试的模型,一般保存准确率最高的

(1)保存模型

这时我们只需要保存模型参数就行了

torch.save(model.state_dict, path)

(2)恢复模型

model.load_state_dict(torch.load(path))

 

以上是关于colab pytorch保存模型的主要内容,如果未能解决你的问题,请参考以下文章

使用谷歌colab训练crnn模型

在 Colab TPU 上运行 Pytorch 堆叠模型

深度学习及pytorch基础

如何将自定义 Pytorch 模型转换为 torchscript(pth 到 pt 模型)?

如何确保所有 PyTorch 代码充分利用 Google Colab 上的 GPU

神经网络学习小记录70——Pytorch 使用Google Colab进行深度学习