PyTorch GPU使用方式

首先通過：

torch.cuda.is_available()

看你的pytorch是否支持CUDA計算，確認支持后：

1.在終端執行程序時設置使用GPU:

CUDA_VISIBLE_DEVICES=1 python train.py

Environment Variable Syntax      Results

CUDA_VISIBLE_DEVICES=1           Only device 1 will be seen
CUDA_VISIBLE_DEVICES=0,1         Devices 0 and 1 will be visible
CUDA_VISIBLE_DEVICES="0,1"       Same as above, quotation marks are optional
CUDA_VISIBLE_DEVICES=0,2,3       Devices 0, 2, 3 will be visible; device 1 is masked
CUDA_VISIBLE_DEVICES=""          No GPU will be visible

 

2.python代碼中設置使用GPU

方法一:

device = torch.device('cuda:0') #數字切換卡號
# device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)
data.to(device)

方法二:

model.cuda(2)
data.cuda(2)

方法三:

with torch.cuda.device(2):
    model.cuda()
    data.cuda()

方法四:

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2" #GPU編號

 

多GPU實驗:

Data Parallelism數據並行是將mini-batch,分割成多個更小的mini-batches的方法,然后並行地計算各個小的mini-batches.

Data Parallelism通過 torch.nn.DataParallel 來實施,可以將一個模塊包裝在DataParallel中，並將在batch維度中通過多個GPU對其進行並行化。

DataParallel

import torch
import torch.nn as nn


class DataParallelModel(nn.Module):

    def __init__(self):
        super().__init__()
        self.block1 = nn.Linear(10, 20)

        # wrap block2 in DataParallel
        self.block2 = nn.Linear(20, 20)
        self.block2 = nn.DataParallel(self.block2)

        self.block3 = nn.Linear(20, 20)

    def forward(self, x):
        x = self.block1(x)
        x = self.block2(x)
        x = self.block3(x)
        return x

-------------------------------------------------------------------------------------------------------------------------------------

CLASS torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0 )

pytorch文檔鏈接

在模塊級別實現數據並行性。

此容器在batch維度上根據device_ids中指定的devices(數量)拆分輸入（其他對象將在每個device上復制一次），從而並行化給定module的應用程序。在forward過程中，module在每個device上被復制，每個副本處理一部分輸入。在backwards過程中，來自每個副本的gradients被匯總到原始module中。

batch size需要比用到的GPUs數量大.

        warning

        在multi-GPU訓練時,推薦使用DistributedDataParallel(分布式數據並行)來代替DataParallel,even if there is only a single node. See: Use nn.parallel.DistributedDataParallel instead of multiprocessing or nn.DataParallel and Distributed Data Parallel.

允許將任意位置和關鍵字輸入傳遞到DataParallel中，但某些類型需要特殊處理。張量將分散在指定的dim上（默認為0)。

# dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs

tuple、list和dict類型將被淺層復制。其他類型將在不同的線程之間共享，如果在模型forward過程中寫入，則可能會損壞.

在運行DataParallel模塊之前，必須將需要並行化的module的參數和緩沖區傳入到device_ids [0]中。

        warning

        在forward過程中,module在每個device上都會被復制,因此，對正在forward的模塊的任何更新都將丟失。舉個例子,如果module在每個forward過程中都有一個遞增的計數器,它會始終保持在初始值，因為更新是在副本上完成的,forward完成后副本會銷毀.但是，DataParallel保證device [0]上的副本具有與基本並行化module共享的參數和緩沖區。所以,會記錄對device [0]上參數或緩沖區的更新。E.g. BatchNorm2d和spectral_norm()依賴此行為來更新緩沖區。

        warning

         forward和backward hooks 被定義在module中,module中的子模塊會被調用 len(device_ids) 次,每個輸入都位於特定device上。特別地，會保證hooks在相對應的device上的操作以正確的順序執行。例如,不能保證通過register_forward_pre_hook()設置的hooks都在len(devices_ids) forward()的調用前執行.但每個此類hook都應在該device的相應forward（）調用之前執行。

        warning

        當module的forward()返回一個標量(0維張量),該包裝器將返回一個向量，該向量的長度等於數據並行性中使用的devices的數量，其中包含每個device的結果。

       Parameters

• module (Module) – module to be parallelized

• device_ids (list of python:int or torch.device) – CUDA devices (default: all devices)

• output_device (int or torch.device) – device location of output (default: device_ids[0])

        Variables

        ~DataParallel.module (Module) – the module to be parallelized

舉例:

net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])
output = net(input_var)  # input_var can be on any device, including CPU

----------------------------------------------------------------------------------------------------------------------------------------------

無需在CPU模式下更改代碼

包裝模塊的屬性

用DataParallel包裝Module后,該module的屬性(例如.自定義的方法)變得不可訪問.這是因為DataParallel定義了一些新的成員,允許訪問其他屬性可能會導致其名稱的沖突.對於一些仍然想要訪問這些不可用屬性的用戶，一種解決方法是使用DataParallel的子類，如下所示。

class MyDataParallel(nn.DataParallel): #DataParallel的子類
    def __getattr__(self, name):
        return getattr(self.module, name)

在其上實現DataParallel的原函數:

一般來說,pytorch的nn.parallel原函數可以單獨使用,我們已經實現了類似於MPI的簡單原函數:

• replicate(復制): 復制Module到多個devices上
• scatter(分散): 在第一個維度(batch維度)上分配數據
• gather(聚合): 在第一維度上收集連接輸入
• parallel_apply(並行應用): 將一組已分配的輸入傳給一組已分配的models.

為了更好的說明,這里使用這些集合組成的data_parallel函數

def data_parallel(module, input, device_ids, output_device=None):
    if not device_ids:
        return module(input)

    if output_device is None:
        output_device = device_ids[0]

    replicas = nn.parallel.replicate(module, device_ids)
    inputs = nn.parallel.scatter(input, device_ids)
    replicas = replicas[:len(inputs)]
    outputs = nn.parallel.parallel_apply(replicas, inputs)
    return nn.parallel.gather(outputs, output_device)

部分模型在CPU上部分模型在GPU上

讓我們看一個實現網絡的小例子，其中網絡的一部分在CPU上，另一部分在GPU上

device = torch.device("cuda:0")

class DistributedModel(nn.Module):

    def __init__(self):
        super().__init__(
            embedding=nn.Embedding(1000, 10),
            rnn=nn.Linear(10, 10).to(device),
        )

    def forward(self, x):
        # Compute embedding on CPU
        x = self.embedding(x)

        # Transfer to GPU
        x = x.to(device)

        # Compute RNN on GPU
        x = self.rnn(x)
        return x