关于torch.load()变了cuda位置还是cuda内存不够的问题
问题背景:在一次任务中,由于需要使用cuda进行代码运行,但是分明修改了cuda到一个空闲的卡位置,但是依然抱错cuda out of memory的问题
在任务中,最开始原始代码是采用以下方式尽心模型load的(指定了cuda to device的位置是空闲的卡2):
但是依然抱错cuda out of memory,分析如下:
model.load_state_dict(torch.load(FLAGS.restore))
print(f"Model loaded from {FLAGS.restore}")
model.to(FLAGS.device)
从上面的代码可以看出,其实model是先load到默认路径,然后再转换到了指定的device的位置,这样就在第一次load的时候,如果默认位置的cuda被占用,就会导致out of memory的问题
建议好的习惯:
model.load_state_dict(torch.load(FLAGS.restore, map_location=‘cuda:2’))
在model load 的时候直接map_location到指定位置,例如你可以在config文件中设置cuda卡配置:
model.load_state_dict(torch.load(FLAGS.restore, map_location=FLAGS.device))
这样则免除了load到默认卡的这一步。
respect!!!