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