在Torch中,可以使用torch.save()
函数来保存模型,使用torch.load()
函数来加载模型。
保存模型示例代码:
import torch
# 定义模型
model = YourModel()
# 保存模型
torch.save(model.state_dict(), 'model.pth')
加载模型示例代码:
import torch
from your_model_module import YourModel
# 创建模型实例
model = YourModel()
# 加载模型参数
model.load_state_dict(torch.load('model.pth'))
model.eval() # 设置为评估模式
辰迅云「云服务器」,即开即用、新一代英特尔至强铂金CPU、三副本存储NVMe SSD云盘,价格低至29元/月。点击查看>>
推荐阅读: 怎么指定优化算法和超参数来训练Torch模型