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