pytorch 模型的加载

刚开始用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'])