pytorch使用指定的GPU

###pytorch使用指定的GPU
在实践中,往往会有一些显卡被其它任务占用,这时我们需要指定使用空闲的GPU进行工作。

通过环境变量指定GPU

  1. export CUDA_VISIBLE_DEVICES=0,1,2,3
  2. 在程序开头设置os.environ[“CUDA_VISIBLE_DEVICES”] = ‘0,1,2,3’
  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)

猜你喜欢

转载自blog.csdn.net/qq_38600065/article/details/106740282