pytorch設置多GPU運行的方法

 

1.DataParallel layers (multi-GPU, distributed)

1）DataParallel

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

實現模塊級別的數據並行

該容器是通過在batch維度上將輸入分到指定的device中來在給定的module應用上實現並行。在前向傳播中，模塊module將在每個設備device上都復制一個，然后每個復制體都會處理一部分的輸入。在后向傳播階段，從每個復制體中計算得到的梯度會加在一起求和，然后傳給最初的模塊module

batch size的大小應該比GPU的個數大

也可見Use nn.DataParallel instead of multiprocessing

任意位置和關鍵字輸入都允許被傳給DataParallel，但是一些類型將進行特定處理。tensors將在指定的維度dim（默認為0）上被分散。tuple、list和dict類型將被淺復制。其他類型將在不同的線程中共享，且在寫入模型的前向傳播時可以被打斷

在運行該DataParallel模塊淺，並行module必須在device_ids[0]有自己的參數和緩沖區

 

警告⚠️

在每次前向傳播中，module在每個設備中復制一遍，所以任何對該正在運行的module的forward的更新都將丟失。比如，如果module有一個計算屬性，在每次forward后自增，那么該屬性將保持在初始狀態，因為已經進行更新的復制體在forward后將被摧毀。但是，DataParallel能保證在device[0]上的復制體的參數和緩沖區與基本並行module共享存儲。因此對參數或device[0]上的緩沖區的內置更新將被保留。比如BatchNorm2d 和 spectral_norm()就依賴於該特性更新緩沖區

 

警告⚠️

定義在module中的前向后向hooks和它的子模塊將被調用len(device_ids)次，每一次都帶着對應設備上的輸入。尤其是hooks僅保證根據相應設備上的操作以正確順序被執行。比如，不保證通過register_forward_pre_hook()被設置的hooks在所有len(device_ids)次forward()調用之前被執行，但是每個這樣的hook將在該設備的相應forward()調用之前被執行

 

警告⚠️

當module在forward()返回scalar（比如0維度的tensor），該封裝器將返回一個有着與使用在數據並行中的設備數量相同的長度向量，包含着來自每個設備的結果

 

注意⚠️

在封裝在DataParallel中的模塊module中使用pack sequence -> recurrent network -> unpack sequence模式是十分微妙的。詳情可見My recurrent network doesn’t work with data parallelism

 

參數：

• module（Module）：並行模塊module
• device_ids(int類型列表或torch.device): CUDA設備（默認為所有設備）
• output_device(int或torch.device):輸出的設備位置（默認為device_ids[0]）

 

Variables：

~DataParallel.module (Module) ：並行模塊

 

例子：

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

 

等價於下面的方法：

2.DataParallel functions (multi-GPU, distributed)

1）data_parallel

torch.nn.parallel.data_parallel(module, inputs, device_ids=None, output_device=None, dim=0, module_kwargs=None)

該給定的device_ids的GPU上並行執行module(input)

這個是上面的DataParallel模塊的函數版本

參數：

• module（Module）：並行模塊module
• input（Tensor）：module的輸入
• device_ids(int類型列表或torch.device): CUDA設備（默認為所有設備）
• output_device(int或torch.device):輸出的設備位置（默認為device_ids[0]）

返回：

包含着位於指定output_device上的model(input)的結果的Tensor

 

例子：

import torch.nn.parallel
...
    def forward(self,input):
        if self.ngpu>1:
             output=nn.parallel.data_parallel(self.model,input,range(self.ngpu)) #在多個gpu上運行模型，並行計算
        else:
            output=self.model(input)
        
        return output