在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() # 設置為評估模式