Сохраните и загрузите контрольный пункт 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.