（原）PyTorch中使用指定的GPU

轉載請注明出處：

http://www.cnblogs.com/darkknightzh/p/6836568.html

 

PyTorch默認使用從0開始的GPU，如果GPU0正在運行程序，需要指定其他GPU。

有如下兩種方法來指定需要使用的GPU。

1. 類似tensorflow指定GPU的方式，使用CUDA_VISIBLE_DEVICES。

1.1 直接終端中設定：

CUDA_VISIBLE_DEVICES=1 python my_script.py

1.2 python代碼中設定：

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

見網址：http://www.cnblogs.com/darkknightzh/p/6591923.html

2. 使用函數 set_device

import torch
torch.cuda.set_device(id)

該函數見 pytorch-master\torch\cuda\__init__.py。

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