pytorch - 有什么方法可以用他们的类保存 pytorch 模型

Question

当我尝试加载我保存的模型时，我需要导入它的类。例如：

from module import Net
torch.load('saved_model.pth')

有什么办法可以避免这种导入？例如用类或其他东西保存模型？

标签： pytorch

Answer 1

如果您想简单地将模型加载到已知的 nn.Module 对象中，例如net可以使用torch.load_state_dict('saved_model.pth'). 如果您想保存整个模型以便其他人可以使用它，您必须使用它来腌制：

import pickle
 
net = Net()
with open('saved_model.pth', 'w') as filehandler: 
    pickle.dump(net, filehandler)

装载：

with open('saved_model.pth', 'w') as filehandler:
    net = pickle.load(filehandler)

但是，强烈建议不要使用 pickle，因为它可能会将自定义内容保存到您的机器/环境中，并导致它无法在其他人的机器上运行。如果你真的必须使用pickle，那么你是否可以将类与神经网络解耦，将类保存在pickle文件中并将参数保存在torch文件中，这可能值得一看。

希望这会有所帮助，而不仅仅是您知道的东西。