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

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

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

1.1 直接終端中設(shè)定：

CUDA_VISIBLE_DEVICES=1 python my_script.py

1.2 python代碼中設(shè)定：

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

見網(wǎng)址：http://www.cnblogs.com/darkknightzh/p/6591923.html

2. 使用函數(shù) set_device

import torch torch.cuda.set_device(id)

該函數(shù)見 pytorch-master\torch\cuda_init_.py。
不過官方建議使用CUDA_VISIBLE_DEVICES，不建議使用 set_device 函數(shù)。

該文章為轉(zhuǎn)載文章，原文章地址:http://www.cnblogs.com/darkknightzh/p/6836568.html

總結(jié)

以上是生活随笔為你收集整理的PyTorch中使用指定的GPU的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。