刚开始用pytorch不久,记录一些遇到的问题
pytorch模型保存
方式一:保存网络结构信息和模型参数信息,保存cnn,命名为“autoencoder.pkl”
torch.save(cnn,'autoencoder.pkl')
方式二:只保存网络训练的模型参数,保存net.state_dict()
torch.save({'epoch':epoch+1, #第几次迭代
'step':step+1, #第几个迭代步
'state_dict':cnn.state_dict()},
'autoencoder.pkl' )
pytorch 模型载入
方式一:如果保存的是网络,载入时只需要加载网络就可以了
cnn = torch.load('autoencoder.pkl')
方式二:如果保存的只是参数,载入时首先需要建立一个跟之前一样的网络,然后得到网络的对象后,在载入模型参数
cnn = Autoencoder()
cnn.load_state_dict(torch.load('autoencoder.pkl')['state_dict'])