1. 利用CUDA_VISIBLE_DEVICES設(shè)置可用顯卡

在CUDA中設(shè)定可用顯卡，一般有2種方式：

(1) 在代碼中直接指定

import os

os.environ['CUDA_VISIBLE_DEVICES'] = gpu_ids

(2) 在命令行中執(zhí)行代碼時(shí)指定

CUDA_VISIBLE_DEVICES=gpu_ids python3 train.py

如果使用sh腳本文件運(yùn)行代碼，則有3種方式可以設(shè)置

(3) 在命令行中執(zhí)行腳本文件時(shí)指定：

CUDA_VISIBLE_DEVICES=gpu_ids sh run.sh

(4) 在sh腳本中指定：

source bashrc

export CUDA_VISIBLE_DEVICES=gpu_ids && python3 train.py

(5) 在sh腳本中指定

source bashrc

CUDA_VISIBLE_DEVICES=gpu_ids python3 train.py

如果同時(shí)使用多個(gè)設(shè)定可用顯卡的指令，比如

source bashrc

export CUDA_VISIBLE_DEVICES=gpu_id1 && CUDA_VISIBLE_DEVICES=gpu_id2 python3 train.py

那么高優(yōu)先級(jí)的指令會(huì)覆蓋第優(yōu)先級(jí)的指令使其失效。優(yōu)先級(jí)順序?yàn)?#xff1a;不使用sh腳本 (1)>(2)； 使用sh腳本(1)>(5)>(4)>(3)

個(gè)人感覺(jué)在煉丹時(shí)建議大家從(2)(3)(4)(5)中選擇一個(gè)指定可用顯卡，不要重復(fù)指定以防造成代碼的混亂。方法(1)雖然優(yōu)先級(jí)最高，但是需要修改源代碼，所以不建議使用。

2 .cuda()方法和torch.cuda.set_device()

我們還可以使用.cuda()[包括model.cuda()/loss.cuda()/tensor.cuda()]方法和torch.cuda.set_device()來(lái)把模型和數(shù)據(jù)加載到對(duì)應(yīng)的gpu上。

(1) .cuda()

以model.cuda()為例，加載方法為：

model.cuda(gpu_id) # gpu_id為int類型變量，只能指定一張顯卡

model.cuda('cuda:'+str(gpu_ids)) #輸入?yún)?shù)為str類型，可指定多張顯卡

model.cuda('cuda:1,2') #指定多張顯卡的一個(gè)示例

(2) torch.cuda.set_device()

使用torch.cuda.set_device()可以更方便地將模型和數(shù)據(jù)加載到對(duì)應(yīng)GPU上, 直接定義模型之前加入一行代碼即可

torch.cuda.set_device(gpu_id) #單卡

torch.cuda.set_device('cuda:'+str(gpu_ids)) #可指定多卡

但是這種寫(xiě)法的優(yōu)先級(jí)低，如果model.cuda()中指定了參數(shù)，那么torch.cuda.set_device()會(huì)失效，而且pytorch的官方文檔中明確說(shuō)明，不建議用戶使用該方法。

第1節(jié)和第2節(jié)所說(shuō)的方法同時(shí)使用是并不會(huì)沖突，而是會(huì)疊加。比如在運(yùn)行代碼時(shí)使用

CUDA_VISIBLE_DEVICES=2,3,4,5 python3 train.py

而在代碼內(nèi)部又指定

model.cuda(1)

loss.cuda(1)

tensor.cuda(1)

那么代碼會(huì)在GPU3上運(yùn)行。原理是CUDA_VISIBLE_DEVICES使得只有GPU2,3,4,5可見(jiàn)，那么這4張顯卡，程序就會(huì)把它們看成GPU0,1,2,3，.cuda(1)把模型/loss/數(shù)據(jù)都加載到了程序所以為的GPU1上，則實(shí)際使用的顯卡是GPU3。

如果利用.cuda()或torch.cuda.set_device()把模型加載到多個(gè)顯卡上，而實(shí)際上只使用一張顯卡運(yùn)行程序的話，那么程序會(huì)把模型加載到第一個(gè)顯卡上，比如如果在代碼中指定了

model.cuda('cuda:2,1')

在運(yùn)行代碼時(shí)使用

CUDA_VISIBLE_DEVICES=2,3,4,5 python3 train.py

這一指令，那么程序最終會(huì)在GPU4上運(yùn)行。

3.多卡數(shù)據(jù)并行torch.nn.DataParallel

多卡數(shù)據(jù)并行一般使用

torch.nn.DataParallel(model,device_ids)

其中model是需要運(yùn)行的模型，device_ids指定部署模型的顯卡，數(shù)據(jù)類型是list

device_ids中的第一個(gè)GPU(即device_ids[0])和model.cuda()或torch.cuda.set_device()中的第一個(gè)GPU序號(hào)應(yīng)保持一致，否則會(huì)報(bào)錯(cuò)。此外如果兩者的第一個(gè)GPU序號(hào)都不是0,比如設(shè)置為：

model=torch.nn.DataParallel(model,device_ids=[2,3])

model.cuda(2)

那么程序可以在GPU2和GPU3上正常運(yùn)行，但是還會(huì)占用GPU0的一部分顯存(大約500M左右)，這是由于pytorch本身的bug導(dǎo)致的(截止1.4.0，沒(méi)有修復(fù)這個(gè)bug)。

device_ids的默認(rèn)值是使用可見(jiàn)的GPU，不設(shè)置model.cuda()或torch.cuda.set_device()等效于設(shè)置了model.cuda(0)

4. 多卡多線程并行torch.nn.parallel.DistributedDataParallel

(這個(gè)我是真的沒(méi)有搞懂,,,,)

參考了這篇文章和這個(gè)代碼，關(guān)于GPU的指定，多卡多線程中有2個(gè)地方需要設(shè)置

torch.cuda.set_device(args.local_rank)

torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])

模型/loss/tensor設(shè)置為.cuda()或.cuda(args.local_rank)均可，不影響正常運(yùn)行。

5. 推薦設(shè)置方式：

(1) 單卡

使用CUDA_VISIBLE_DEVICES指定GPU，不要使用torch.cuda.set_device()，不要給.cuda()賦值。

(2) 多卡數(shù)據(jù)并行

直接指定CUDA_VISIBLE_DEVICES，通過(guò)調(diào)整可見(jiàn)顯卡的順序指定加載模型對(duì)應(yīng)的GPU,不要使用torch.cuda.set_device()，不要給.cuda()賦值，不要給torch.nn.DataParallel中的device_ids賦值。比如想在GPU1,2,3中運(yùn)行，其中GPU2是存放模型的顯卡，那么直接設(shè)置

CUDA_VISIBLE_DEVICES=2,1,3

(3) 多卡多線程

由于這塊還有好多地方?jīng)]有搞懂(尤其是torch.nn.parallel.DistributedDataParallel)，所以文章中難免會(huì)有很多錯(cuò)誤和疏漏，歡迎各位大佬指正。

總結(jié)

以上是生活随笔為你收集整理的pytorch 指定卡1_在pytorch中指定显卡的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。

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