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