torch.save(model, 'model.pth')
model = torch.load('model.pth')
torch.save(model.state_dict(), 'model.pth')
model.load_state_dict(torch.load('model.pth')
model.load_state_dict(torch.load('model.pth', map_location='cpu'))
- 假设我们只保存了模型的参数
model.state_dict()
到文件名为model_state_dict.pth, model = Net()
- 1. cpu -> cpu或者gpu -> gpu:
checkpoint = torch.load('model_state_dict.pth')
model.load_state_dict(checkpoint)
torch.load('model_state_dict.pth', map_location=lambda storage, loc: storage.cuda(1))
torch.load('model_state_dict.pth', map_location={'cuda:1':'cuda:0'})
torch.load('model_state_dict.pth', map_location=lambda storage, loc: storage)