保存模型:
torch.save(mlp1, r'XXXX\net1.pkl') # 保存整个神经网络到net1.pkl中
torch.save(mlp1.state_dict(), r'XXXX\net1_params.pkl') # 保存网络里的参数
读取模型:
net = torch.load(r'XXXX\net1.pkl')
net.load_state_dict(torch.load('XXX/net1_params.pkl'))
或许会报错:
AttributeError: Can't get attribute 'XXX1' on <module '__main__' from 'XXX2.py'>
原因是:之前生成net1的py文件代码中定义了XXX1类,现在调用模型的XXX2.py文件中没有这个类,只需要将这个类复制过来就行。