PyTorch中如何使用指定的GPU

PyTorch默认使用从0开始的GPU,如果GPU0正在运行程序,需要指定其他GPU。

hero.png


有如下两种方法来指定需要使用的GPU。
  • 使用CUDA_VISIBLE_DEVICES,类似tensorflow指定GPU的方法。
  • 直接终端中设定:
    CUDA_VISIBLE_DEVICES=1 python my_script.py
  • python代码中设定:
    import osos.environ["CUDA_VISIBLE_DEVICES"] = "2"
  • 使用函数 set_device,该函数见 pytorch-master\torch\cuda\__init__.py。
    import torchtorch.cuda.set_device(id)

 
不过官方建议使用CUDA_VISIBLE_DEVICES,不建议使用 set_device 函数。

0 个评论

要回复文章请先登录注册