Skip to content

GitLab

Commit 39ad741a authored by Yizhou Wang's avatar Yizhou Wang
Browse files

save loss_ave to checkpoint

parent 8568630b
Show whitespace changes
Inline Side-by-side
Showing with 6 additions and 2 deletions
tools/train.py
View file @ 39ad741a
...@@ -186,6 +186,8 @@ if __name__ == "__main__": ...@@ -186,6 +186,8 @@ if __name__ == "__main__":
loss_cp = checkpoint['loss'] loss_cp = checkpoint['loss']
if 'iter_count' in checkpoint: if 'iter_count' in checkpoint:
iter_count = checkpoint['iter_count'] iter_count = checkpoint['iter_count']
if 'loss_ave' in checkpoint:
loss_ave = checkpoint['loss_ave']
else: else:
rodnet.load_state_dict(checkpoint) rodnet.load_state_dict(checkpoint)
...@@ -281,7 +283,8 @@ if __name__ == "__main__": ...@@ -281,7 +283,8 @@ if __name__ == "__main__":
'iter': iter, 'iter': iter,
'model_state_dict': rodnet.state_dict(), 'model_state_dict': rodnet.state_dict(),
'optimizer_state_dict': optimizer.state_dict(), 'optimizer_state_dict': optimizer.state_dict(),
'loss': loss_confmap, 'loss': loss_confmap.item(),
'loss_ave': loss_ave,
'iter_count': iter_count, 'iter_count': iter_count,
} }
save_model_path = '%s/epoch_%02d_iter_%010d.pkl' % (model_dir, epoch + 1, iter_count + 1) save_model_path = '%s/epoch_%02d_iter_%010d.pkl' % (model_dir, epoch + 1, iter_count + 1)
...@@ -298,7 +301,8 @@ if __name__ == "__main__": ...@@ -298,7 +301,8 @@ if __name__ == "__main__":
'iter': iter, 'iter': iter,
'model_state_dict': rodnet.state_dict(), 'model_state_dict': rodnet.state_dict(),
'optimizer_state_dict': optimizer.state_dict(), 'optimizer_state_dict': optimizer.state_dict(),
'loss': loss_confmap, 'loss': loss_confmap.item(),
'loss_ave': loss_ave,
'iter_count': iter_count, 'iter_count': iter_count,
} }
save_model_path = '%s/epoch_%02d_final.pkl' % (model_dir, epoch + 1) save_model_path = '%s/epoch_%02d_final.pkl' % (model_dir, epoch + 1)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment