在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() # 设置为评估模式