###pytorch使用指定的GPU
在实践中,往往会有一些显卡被其它任务占用,这时我们需要指定使用空闲的GPU进行工作。
通过环境变量指定GPU
- export CUDA_VISIBLE_DEVICES=0,1,2,3
- 在程序开头设置os.environ[“CUDA_VISIBLE_DEVICES”] = ‘0,1,2,3’
- 在命令行执行的命令前添加CUDA_VISIBLE_DEVICES=0,1,2,3,4
使用torch.cuda接口
torch.cuda.set_device(0)
亲测貌似不行,还是会报cuda: out of memory错误。知道为什么的朋友可以留言告诉我。
在模型初始化时指定GPU
net = Net.cuda(0)