Skip to content

Instantly share code, notes, and snippets.

@giacaglia
Created December 9, 2019 01:28
Show Gist options
  • Save giacaglia/aa74fc895fb62f8338ad752cbb36737b to your computer and use it in GitHub Desktop.
Save giacaglia/aa74fc895fb62f8338ad752cbb36737b to your computer and use it in GitHub Desktop.
# Save checkpoint
checkpoint = {
'model': model.state_dict(),
'optimizer': optimizer.state_dict(),
'amp': amp.state_dict()
}
torch.save(checkpoint, 'amp_checkpoint.pt')
...
# Restore
model = ...
optimizer = ...
checkpoint = torch.load('amp_checkpoint.pt')
model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
amp.load_state_dict(checkpoint['amp'])
# Continue training
...
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment