Pytorch FAQ

1、pytorch指定使用哪個GPU

使用CUDA_VISIBLE_DEVICES，類似tensorflow指定GPU的方法。 直接終端中設定：

CUDA_VISIBLE_DEVICES=1 python my_script.py

python代碼中設定：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"

使用函數set_device，該函數見 pytorch-master\torch\cuda\__init__.py。

import torchtorch.cuda.set_device(id)

不過官方建議使用CUDA_VISIBLE_DEVICES，不建議使用set_device函數。