torch 保存模型

保存

state = {
    
    'net': model.state_dict(), 'optimizer':optimizer.state_dict(), 'epoch': epoch}
torch.save(state, 'dir.pth')        

加载

checkpoint = torch.load(dir)
model.load_state_dict(checkpoint['net'])
optimizer.load_state_dict(checkpoint['optimizer'])
start_epoch = checkpoint['epoch'] + 1