os.environ[CUDA_VISIBLE_DEVICES] 失效无法指定 GPU

问题描述：

用语句 os.environ['CUDA_VISIBLE_DEVICES'] 设置了程序跑的GPU，但是模型还是跑在了默认的第一块卡——‘0’卡上。

---

原因分析：

网上很多分析都说¹：

os.environ[‘CUDA_VISIBLE_DEVICES’] 必须在import torch之前。

但是亲测并不是，例如以下案例：

import torch
import os

def main():
	...
    os.environ['CUDA_VISIBLE_DEVICES'] = 1
    ...
    ...

if __name__ == '__main__':
    main()

实际运行，还是可以设置运行在哪块GPU的。

查阅资料发现，关键原因在于：²

you need to set that before the first use of cuda rather than after that

即设置GPU需要在第一次使用 CUDA 之前才行。

---

解决方案：

通过定位发现，本人出现问题的语句在于开头导入的包

from vector_quantize_pytorch import VectorQuantize

查找这个包的源码可知³，此包调用时还使用了 CUDA ：

import torch  # 无影响
from torch import nn, einsum
import torch.nn.functional as F
import torch.distributed as distributed
from torch.cuda.amp import autocast

因此，解决方案是在导入 vector_quantize_pytorch 之前设置GPU。

---

总结

1. 为了避免不必要的麻烦，有些博客说的方法是对的⁴，即 在import之前配置好GPU。（注意，它们所说的需要在import torch之前设置其实实测不影响。）
2. 在命令行或者.sh文件里，在运行.py文件前直接指定GPU⁵：CUDA_VISIBLE_DEVICES=x python xxx.py，其他方法可参考Pytorch 指定GPU。

---

参考资料

---

1. os.environ[‘CUDA_VISIBLE_DEVICES‘] 无法生效原因 ↩︎

2. os.environ[‘CUDA_VISIBLE_DEVICES’]不生效的原因 | 不管是tf还是pytorch ↩︎

3. vector_quantize_pytorch.py ↩︎

4. os.environ[“CUDA_VISIBLE_DEVICES“]设置GPU不成功 ↩︎

5. Pytorch 指定GPU ↩︎