如何在Torch中保存和加載模型

小樊
64
2024-03-30 19:33:03

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

保存模型：

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

加載模型：

# 加載模型
model = Model()
model.load_state_dict(torch.load('model.pth'))
model.eval()

在加載模型時(shí)，要確保模型結(jié)構(gòu)和保存時(shí)的結(jié)構(gòu)一致，否則可能會(huì)出現(xiàn)錯(cuò)誤。

0 踩