os.environ[‘CUDA_VISIBLE_DEVICES‘]无效的解决方法

文章首发及后续更新：https://mwhls.top/3989.html，无图/无目录/格式错误/更多相关请至首发页查看。
新的更新内容请到mwhls.top查看。
欢迎提出任何疑问及批评，非常感谢！

场景

• 需要指定 GPU 1 来进行训练，但 os.environ['CUDA_VISIBLE_DEVICES'] = '1' 无效

原因

• 在 torch 引入前指定 GPU 才有效

解决方法

• 创建 startup.py，在 startup.py 中先指定GPU，然后再 from tools/train import main
  • 其中，main 指的是训练主函数

示例

• 代码顺序执行，进入 startup()，

  1. 然后会先设置 GPU

  2. 再引入 main()，

  3. 再进入 parse_args() 引入 DictAction，

  4. 再进入 main() 进行训练

• 这两个放在 os.environ[] 后 import 的都会引入 torch。

import os
import argparse


def parse_args(config1, config2):
    from mmcv.utils import DictAction				# 会import torch
    pass

def startup(gpu=0, config1='xxx', config2='xxx'):
    os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)
    from tools.train import main					# 会import torch
    args = parse_args(config1, config2)
    main(args)
    
if __name__ == "__main__":
    startup(gpu=1)

参考文献

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