參考了官方文檔， torch.nn.parallel.DataParallel
以及https://zhuanlan.zhihu.com/p/102697821

在運(yùn)行此DataParallel模塊之前，并行化模塊必須在device_ids [0]上具有其參數(shù)和緩沖區(qū)。在執(zhí)行DataParallel之前，會(huì)首先把其模型的參數(shù)放在device_ids[0]上，一看好像也沒有什么毛病，其實(shí)有個(gè)小坑。我舉個(gè)例子，服務(wù)器是八卡的服務(wù)器，剛好前面序號(hào)是0的卡被別人占用著，于是你只能用其他的卡來，比如你用2和3號(hào)卡，如果你直接指定device_ids=[2, 3]的話會(huì)出現(xiàn)模型初始化錯(cuò)誤，類似于module沒有復(fù)制到在device_ids[0]上去。

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = "2, 3"

當(dāng)你添加這兩行代碼后，那么device_ids[0]默認(rèn)的就是第2號(hào)卡，你的模型也會(huì)初始化在第2號(hào)卡上了，而不會(huì)占用第0號(hào)卡了。這里簡單說一下設(shè)置上面兩行代碼后，那么對(duì)這個(gè)程序而言可見的只有2和3號(hào)卡，和其他的卡沒有關(guān)系，這是物理上的號(hào)卡，邏輯上來說其實(shí)是對(duì)應(yīng)0和1號(hào)卡，即device_ids[0]對(duì)應(yīng)的就是第2號(hào)卡，device_ids[1]對(duì)應(yīng)的就是第3號(hào)卡。（當(dāng)然你要保證上面這兩行代碼需要定義在

device_ids = [0, 1] net = torch.nn.DataParallel(net, device_ids=device_ids)

而且不要以為添加了to(device)或者.cuda（） 就運(yùn)行在卡0上了，其實(shí)不然！ 這只是將模型暫時(shí)放到了卡0上，第二步系統(tǒng)會(huì)復(fù)制模型和參數(shù)的！

而且并行和to、cuda是缺一不可的！

代碼測試：

import torch import torch.nn as nn import torch.nn.functional as F import osos.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" #指定程序可見的devices！ 也就是其它設(shè)備都不可見 os.environ["CUDA_VISIBLE_DEVICE"] = "0, 1" #設(shè)置當(dāng)前使用的GPU設(shè)備為1,0號(hào)兩個(gè)設(shè)備,名稱依次為'/gpu:0'、'/gpu:1'。表示優(yōu)先使用1號(hào)設(shè)備,然后使用0號(hào)設(shè)備## 看起來沒事，但是有坑的！ 如果我們和別人一塊使用的話，別人在用卡0，因?yàn)槌绦驎?huì)默認(rèn)將模型參數(shù)和緩存放到卡0，那么你不改上面的話就會(huì)出現(xiàn)問題device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")class Model(nn.Module):def __init__(self):super(Model, self).__init__()self.conv1 = nn.Conv2d(1, 20, 5)self.conv2 = nn.Conv2d(20, 20, 5)def forward(self, x):x = F.relu(self.conv1(x))return F.relu(self.conv2(x))input = torch.randn(128, 1, 256, 128) net = Model() if torch.cuda.device_count() > 1:# model.to(device)device_ids = [0, 1, 2]net = torch.nn.parallel.DataParallel(net, device_ids=device_ids)net.cuda()# net.to(device) //兩種方法都可以！ 只不是這種方法考慮了不能用的情況！for param in next(net.parameters()):print(param, param.device)print(next(net.parameters()).device) net = net(input) print(net)

總結(jié)

以上是生活随笔為你收集整理的如何并行运行程序的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

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