python-3.x - 保存和加载检查点 pytorch

Question

我制作一个模型并将配置保存为：

def checkpoint(state, ep, filename='./Risultati/checkpoint.pth'):  
    if ep == (n_epoch-1):
        print('Saving state...')
        torch.save(state,filename)
checkpoint({'state_dict':rnn.state_dict()},ep) 

然后我想加载这个配置：

state_dict= torch.load('./Risultati/checkpoint.pth')
    rnn.state_dict(state_dict)

当我尝试时，这是错误：

Traceback (most recent call last):
File "train.py", line 288, in <module>
rnn.state_dict(state_dict)
File "/home/marco/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 593, in state_dict
destination._metadata[prefix[:-1]] = dict(version=self._version)
AttributeError: 'dict' object has no attribute '_metadata'

我哪里做错了？

提前谢谢

Answer 1

您需要加载rnn.state_dict()存储在您加载的字典中：

rnn.load_state_dict(state_dict['state_dict'])

查看load_state_dict方法以获取更多信息。