不管怎么指定pytorch都使用gpu0进行训练解决方法

这个问题困惑了我好久,不管怎么指定,程序中用的
os.environ['CUDA_VISIBLE_DEVICES'] = '1,2'
并不管用
最后都显示

RuntimeError: CUDA out of memory. Tried to allocate 158.00 MiB (GPU 0; 11.17 GiB total capacity; 556.72 MiB already allocated; 26.06 MiB free; 17.28 MiB cached)

真的是很泪奔了,最后查到,pytorch在初始化的时候会默认在第0块显卡上进行,会占用一定的显存,这就导致,在第0块显卡空闲内存不多时,程序会反复报上面的错误,解决方法是使用with语句,在with块下再写pytorch的调用,亲测有效。

if __name__ == '__main__':
  with torch.cuda.device(7):
    opt = opts().parse()
    main(opt)

问题解决。

参考:https://blog.csdn.net/weixin_38376691/article/details/96435895?utm_medium=distribute.pc_relevant.none-task-blog-BlogCommendFromMachineLearnPai2-2.nonecase&depth_1-utm_source=distribute.pc_relevant.none-task-blog-BlogCommendFromMachineLearnPai2-2.nonecase

©️2020 CSDN 皮肤主题: 大白 设计师:CSDN官方博客 返回首页