TrajectoryNet icon indicating copy to clipboard operation
TrajectoryNet copied to clipboard

Published 20 hours ago • verified publisher icon KrishnaswamyLab

Does not use GPU support on Colab

Open SK124 opened this issue 5 years ago • 5 comments

The model is not at all utilising the GPU support Google colab offers. Any changes I should make to the code and in which line would be helpful. @atong01

SK124 avatar Oct 04 '20 06:10 SK124

By default it is set up to use a gpu if available. Please check that one is available with torch.cuda.isavailable(). There is a --gpu flag for specifying gpu and a --use_cpu for the obvious.

atong01 avatar Oct 04 '20 13:10 atong01

Thanks I will look into that btw, i noticed a comment which specified that few operations might require torch<1.5 is that resolved?

SK124 avatar Oct 14 '20 06:10 SK124

Things were broken in pytorch 1.4 specifically for the density regularization on GPU anything 1.5 and up should be fine.

atong01 avatar Oct 14 '20 14:10 atong01

Is there any way i can start training from a checkpoint rather than start from scratch. which line of code should I change?

SK124 avatar Oct 14 '20 15:10 SK124

you should be able to load weights from a checkpoint like this, maybe inserting this around line 450 in main.py (and obviously with a correct path) would do the trick

state_dict = torch.load(args.save + "/checkpt.pth", map_location=device)
model.load_state_dict(state_dict["state_dict"])

atong01 avatar Oct 14 '20 19:10 atong01