个人技术分享

一、保存整个模型

这种方法会保存模型的整个定义(包括网络结构和参数)。这样保存的模型可以直接被加载并使用,但前提是保存和加载模型的代码环境需要完全一致。

保存模型时使用代码:

torch.save(model, 'model.pth')

加载模型时使用代码:

model = torch.load('model.pth')

二、只保存模型参数

这种方法只会保存模型的状态字典(state_dict),即模型的参数。加载这样的模型需要先定义模型的结构,然后加载参数。

保存模型使用代码:

torch.save(model.state_dict(), 'params.pth')

加载模型使用代码:

model = Model()  # 需要先定义模型结构
model.load_state_dict(torch.load('params.pth'))