溫馨提示×

Torch中如何進行模型保存和加載

小樊
108
2024-03-25 11:06:55

在Torch中，可以使用torch.save()函數(shù)來保存模型，使用torch.load()函數(shù)來加載模型。

保存模型示例代碼：

import torch

# 定義模型
model = YourModel()

# 保存模型
torch.save(model.state_dict(), 'model.pth')

加載模型示例代碼：

import torch
from your_model_module import YourModel

# 創(chuàng)建模型實例
model = YourModel()

# 加載模型參數(shù)
model.load_state_dict(torch.load('model.pth'))
model.eval()  # 設(shè)置為評估模式

0 踩