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