介绍Pytorch的一些使用方法

GPU

查看 GPU 信息

更多接口，参考 torch.cuda

torch.cuda.is_available()       # 判断 GPU 是否可用
torch.cuda.device_count()       # 判断有多少 GPU
torch.cuda.get_device_name(0)   # 返回 gpu 名字，设备索引默认从 0 开始
torch.cuda.current_device()     # 返回当前设备索引

torch.device

torch.device 表示 torch.Tensor 分配到的设备的对象。其包含一个设备类型（cpu 或 cuda），以及可选的设备序号。如果设备序号不存在，则为当前设备，即 torch.cuda.current_device() 的返回结果。

可以通过如下方式创建 torch.device 对象：

# 通过字符串
device = torch.device('cpu')
device = torch.device('cuda:1')  # 指定类型及编号。注意，代码不会检查编号是否合法
device = torch.device('cuda')    # 默认为当前设备，如果是多GPU，默认使用全部GPU

还可以通过设备类型加上编号，来创建 device 对象：

1 2	device = torch.device('cuda', 0) device = torch.device('cpu', 0)

配置 CUDA 访问限制

可以通过如下方式，设置当前 Python 脚本可见的 GPU。

在命令行设置

1	CUDA_VISIBLE_DEVICES=1 python my_script.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

在 Python 代码中设置

1 2	import os os.environ["CUDA_VISIBLE_DEVICES"] = "0, 2"

使用函数 set_device

1 2	import torch torch.cuda.set_device(id)

官方建议使用 CUDA_VISIBLE_DEVICES，不建议使用 set_device 函数。

用 GPU 训练

默认情况下，使用 CPU 训练模型。可以通过如下方式，通过 GPU 进行训练。使用 GPU 时，模型和输入必须位于同一张 GPU 上。

.to(device) 和 .cuda() 的区别如下：

to和cuda的区别
to和cuda区别

.to() 中的参数必不可少
对于 module 而言，.to() 是 inplace 的，而 .cuda() 不是；而对于 tensor 而言，两者一致。

注：实测，两者时间消耗持平。推荐使用.to()函数

方式 1 ：使用cuda

device = torch.device("cuda:1")   # 指定模型训练所在 GPU

# 将 CPU 转移至 GPU
if torch.cuda.is_available() and use_gpu:
    net = net.cuda(device)    # 默认在第一块 GPU 上训练

# 同时将数据转移至 GPU,注意cuda有返回值
if torch.cuda.is_available() and use_gpu:
    inputs = inputs.cuda(device)
    labels = labels.cuda(device)

方法 2 ：使用to

device = torch.device("cuda:1")   # 指定模型训练所在 GPU

# 将 CPU 转移至 GPU
if torch.cuda.is_available() and use_gpu:
    net = net.to(device)    # 默认在第一块 GPU 上训练

# 同时将数据转移至 GPU
if torch.cuda.is_available() and use_gpu:
    inputs = inputs.to(device)
    labels = labels.to(device)

存在的问题

batch size 太大

当想要用大批量进行训练，但是 GPU 资源有限，此时可以通过梯度累加（accumulating gradients）的方式进行。

梯度累加的基本思想在于，在优化器更新参数前，也就是执行 optimizer.step() 前，进行多次反向传播，使得梯度累计值自动保存在 parameter.grad 中，最后使用累加的梯度进行参数更新。

这个在 PyTorch 中特别容易实现，因为 PyTorch 中，梯度值本身会保留，除非我们调用 model.zero_grad() 或 optimizer.zero_grad()。

修改后的代码如下所示：

model.zero_grad() # 重置保存梯度值的张量

for i, (inputs, labels) in enumerate(training_set):
    predictions = model(inputs)# 前向计算
    loss = loss_function(predictions, labels)# 计算损失函数
    loss.backward()# 计算梯度
    if (i + 1) % accumulation_steps == 0:#重复多次前面的过程
        optimizer.step()#更新梯度
        model.zero_grad()#重置梯度

model 太大

当模型本身太大，以至于不能放置于一个 GPU 中时，可以通过梯度检查点 (gradient-checkpoingting) 的方式进行处理。

梯度检查点的基本思想是以计算换内存。具体来说就是，在反向传播的过程中，把梯度切分成几部分，分别对网络上的部分参数进行更新。如下图所示：

梯度检查点图示

这种方法速度很慢，但在某些例子上很有用，比如训练长序列的 RNN 模型等。

具体可参考：From zero to research — An introduction to Meta-learning

单机多卡训练，即并行训练。并行训练又分为数据并行 (Data Parallelism) 和模型并行两种。

数据并行指的是，多张 GPU 使用相同的模型副本，但是使用不同的数据批进行训练。而模型并行指的是，多张GPU 分别训练模型的不同部分，使用同一批数据。

两者对比如下图所示：

模型并行 VS 数据并行

数据并行

Pytorch多GPU官方实例

Pytorch API

【Class 原型】

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

【参数】

module ：要进行并行的 module。这里隐含了一点，即网络中的某一层也是可以进行数据并行的，但是一般不会这么使用。
device_ids : CUDA 列表，可以为 torch.device 类型，也可以是编号组成的 int 列表。默认使用全部 GPU
output_device : 某一 GPU 编号或 torch.device 。指定输出的 GPU，默认为第一个，即 device_ids[0]

【返回值】

要进行并行的模型。

【基本使用方式】

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

数据并行的原理

数据并行的具体原理流程为：

将模型加载至主设备上，作为 controller，一般设置为 cuda:0
在每次迭代时，执行如下操作：
1. 将 controller 模型复制（broadcast）到每一个指定的 GPU 上
2. 将总输入的数据 batch，进行均分，分别作为各对应副本的输入 (scatter)
3. 每个副本独立进行前向传播，并进行反向传播，但只是求取梯度，每个GPU上的loss都要进行loss.backward(),得到各自的梯度
4. 将各副本的梯度汇总（gather）到 controller 设备，并进行求和 (reduced add)
  
  During the backwards pass, gradients from each replica are summed into the original module.
5. 更具总体度，更新 controller 设备上的参数

注意事项

【警告 1】

设置的 batch size 为总的批量尺寸，其必须大于 GPU 数量。
在 parallelized module 运行之前，必须保证其在 controller 设备上，存在参数和 buffers。
并行的 GPU 列表中，必须包含主 GPU
当 forward() 中，module 返回一个标量，那么并行的结果将返回一个 vector，其长度等于 device 的数量，对应于各个设备的结果。

【警告 2】

在每次前向传播过程中，module 都先会被复制到每一个 device 上。因此，在前向传播中，任何对该运行的 module 的副本的更新，在此后都将会丢失。

比方说，如果 module 有一个 counter 属性，每次前向传播都会进行累加，则它将会保持为初始值。因为更新是发生在模型的副本（在其他 device 上的副本）上的，并且这些更新在前向传播结束之后将会被销毁。

然而，DataParallel 保证 controller 设备上的副本的参数和 buffers 与其他并行的 modules 之间共享存储。因此，如若对 controller device 的参数和 buffers 的更改，将会被记录。例如，BatchNorm2d 和 spectral_norm() 依赖于这种行为来更新 buffers。

【警告 3】

定义于 module 及其子 module 上的前向传播和反向传播 hooks，将会被调用 len(device_ids) 次，每个设备对应一次。

具体来说，hooks 只能保证按照正确的顺序执行对应设备上的操作，即在对应设备上的 forward() 调用之前执行，但是不能保证，在所有 forward)() 执行之前，通过 register_forward_pre_hook() 执行完成所有的 hooks。

【警告 4】

任何位置和关键字 (positional and keyword) 输入都可以传递给 DataParallel，处理一些需要特殊处理的类型。

tensors 将会在指定维度（默认为 0）上被 scattered。 tuple， list 和 dict 类型则会被浅拷贝。其他类型则会在不同的线程之间进行共享，且在模型前向传播过程中，如果进行写入，则可被打断。

【警告 5】

当对 pack sequence -> recurrent network -> unpack sequence 模式的 module 使用 DataParallel 或 data_parallel 时，有一些小的问题。

每个设备上的 forward 的对应输入，将仅仅是整个输入的一部分。因为默认的 unpack 操作 torch.nn.utils.rnn.pad_packed_sequence() 只会将该设备上的输入 padding 成该设备上的最长的输入长度，因此，将所有设备的结构进行汇总时，可能会发生长度的不匹配的情况。

因此，可以利用 pad_packed_sequence() 的 total_length 参数来保证 forward() 调用返回的序列长度一致。代码如下所示：

from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

class MyModule(nn.Module):
    # ... __init__, other methods, etc.

    # padded_input is of shape [B x T x *] (batch_first mode) and contains
    # the sequences sorted by lengths
    #   B is the batch size
    #   T is max sequence length
    def forward(self, padded_input, input_lengths):
        total_length = padded_input.size(1)  # get the max sequence length
        packed_input = pack_padded_sequence(padded_input, input_lengths,
                                            batch_first=True)
        packed_output, _ = self.my_lstm(packed_input)
        output, _ = pad_packed_sequence(packed_output, batch_first=True,
                                        total_length=total_length)
        return output

m = MyModule().cuda()        # 设置 controller 模型
dp_m = nn.DataParallel(m)    # 进行副本拷贝

示例程序

下面是使用 DataParrel 的核心代码，其余部分与一般的训练流程一致。

# 设置当前脚本可见的 GPU 列表
# 这里设置 0 号和 1 号 GPU 对当前脚本可见。
# 此时，若 DataParallel 中指定使用其他 GPU 资源，额外的编号将会被忽略
os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1"

# 使用数据并行
# 1. 将 model 转移到某 GPU 上 -- net.cuda()
# 2. 指定并行训练要用到的 GPU -- device_ids=[0, 1]
if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    net = nn.DataParallel(net.cuda(), device_ids=[0, 1])

# 将数据转移到 controller 所在 GPU
if torch.cuda.is_available() and use_gpu:
    inputs = inputs.cuda(device)
    labels = labels.cuda(device)

模型的加载

def _Single2Parallel(self, origin_state):
    """
    将串行的权值参数转换为并行的权值参数
    :param origin_state : 原始串行权值参数
    :return             : 并行的权值参数
    """
  	converted = OrderedDict()
  
    for k, v in origin_state.items():
      name = "module." + k
      converted[name] = v
  
  	return converted


def _Parallel2Single(self, origin_state):
  	"""
  	将并行的权值参数转换为串行的权值参数
  	:param origin_state : 原始串行权值参数
  	:return             : 并行的权值参数
  	"""
    
    converted = OrderedDict()
    
    for k, v in origin_state.items():
      name = k[7:]
      converted[name] = v
      
    return converted

模型并行

如果模型本身较大，一张 GPU 放置不下时，要通过模型并行来处理。模型并行指的是，将模型的不同部分，分别放置于不同的 GPU 上，并将中间结果在 GPU 之间进行传递。

尽管从执行时间上来看，将模型的不同部分部署在不同设备上确实有好处，但是它通常是出于避免内存限制才使用。具有特别多参数的模型会受益于这种并行策略，因为这类模型需要很高的内存占用，很难适应到单个系统。

基本使用

下面，我们以一个 toy 模型为例，讲解模型并行。模型并行的实现方式如下所示：

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()

        self.features_1 = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3),
            nn.BatchNorm2d(16),
            nn.ReLU(inplace=True),  # 30
            ......
            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),  # 12
        ).to('cuda:0')

        self.features_2 = nn.Sequential(
            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=2),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),  # 5
            ......).to('cuda:1')  # 1

        self.classifier = nn.Sequential(
            nn.Dropout(),
            ......
            nn.Linear(1024, class_num)).to('cuda:1')

    def forward(self, x):
        out = self.features_1(x.to('cuda:0'))
        out = self.features_2(out.to('cuda:1'))
        out = out.view(-1, 384)
        out = self.classifier(out)
        out = F.softmax(out, dim=1)
        return out

上面的 toy 模型看起来和在单个 GPU 上运行的模型没什么区别，只不过用 to(device) 来将模型内的不同层分散到不同的 GPU 上进行运行，并且将中间结果转移到对应的 GPU 上即可。

backward() 和 torch.optim 将会自动考虑梯度，与在一个 GPU 上没有区别。

注意：在调用 loss 函数时，labels 与 output 必须在同一个 GPU 上。

# 此时，不在此需要使用 model = model.cuda()
model = ToyModel()

loss_fn = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)

optimizer.zero_grad()

for data in trainloader:
    images, labels = data
    
    # 要处理的部分
    images = images.to('cuda:0')
    labels = labels.to('cuda:1')   # 必须与输出所在 GPU 一致
    
    outputs = net(images)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()

模型并行的性能分析

以上的实现解决了单个模型太大，不能存放于一个 GPU 的情况。然而，需要注意的是，相较于在单个 GPU 上运行，其速度更慢。因为任何时候，只有一个 GPU 在工作，而另一个则闲置。而当中间结果在 GPU 之间进行转移时，速度会进一步下降。

下面同时实例分析。以 resnet50 为例，用随机生成的数据输入，比较两个版本的运行时间。

from torchvision.models.resnet import ResNet, Bottleneck

num_classes = 1000


class ModelParallelResNet50(ResNet):
    def __init__(self, *args, **kwargs):
        super(ModelParallelResNet50, self).__init__(
            Bottleneck, [3, 4, 6, 3], num_classes=num_classes, *args, **kwargs)

        self.seq1 = nn.Sequential(
            self.conv1,
            self.bn1,
            self.relu,
            self.maxpool,

            self.layer1,
            self.layer2
        ).to('cuda:0')

        self.seq2 = nn.Sequential(
            self.layer3,
            self.layer4,
            self.avgpool,
        ).to('cuda:1')

        self.fc.to('cuda:1')

    def forward(self, x):
        x = self.seq2(self.seq1(x).to('cuda:1'))
        return self.fc(x.view(x.size(0), -1))

import torchvision.models as models

num_batches = 3
batch_size = 120
image_w = 128
image_h = 128


def train(model):
    model.train(True)
    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(model.parameters(), lr=0.001)

    one_hot_indices = torch.LongTensor(batch_size) \
                           .random_(0, num_classes) \
                           .view(batch_size, 1)

    for _ in range(num_batches):
        # generate random inputs and labels
        inputs = torch.randn(batch_size, 3, image_w, image_h)
        labels = torch.zeros(batch_size, num_classes) \
                      .scatter_(1, one_hot_indices, 1)

        # run forward pass
        optimizer.zero_grad()
        outputs = model(inputs.to('cuda:0'))

        # run backward pass
        labels = labels.to(outputs.device)
        loss_fn(outputs, labels).backward()
        optimizer.step()

import matplotlib.pyplot as plt
plt.switch_backend('Agg')
import numpy as np
import timeit

num_repeat = 10

stmt = "train(model)"

setup = "model = ModelParallelResNet50()"
# globals arg is only available in Python 3. In Python 2, use the following
# import __builtin__
# __builtin__.__dict__.update(locals())
mp_run_times = timeit.repeat(
    stmt, setup, number=1, repeat=num_repeat, globals=globals())
mp_mean, mp_std = np.mean(mp_run_times), np.std(mp_run_times)

setup = "import torchvision.models as models;" + \
        "model = models.resnet50(num_classes=num_classes).to('cuda:0')"
rn_run_times = timeit.repeat(
    stmt, setup, number=1, repeat=num_repeat, globals=globals())
rn_mean, rn_std = np.mean(rn_run_times), np.std(rn_run_times)


def plot(means, stds, labels, fig_name):
    fig, ax = plt.subplots()
    ax.bar(np.arange(len(means)), means, yerr=stds,
           align='center', alpha=0.5, ecolor='red', capsize=10, width=0.6)
    ax.set_ylabel('ResNet50 Execution Time (Second)')
    ax.set_xticks(np.arange(len(means)))
    ax.set_xticklabels(labels)
    ax.yaxis.grid(True)
    plt.tight_layout()
    plt.savefig(fig_name)
    plt.close(fig)


plot([mp_mean, rn_mean],
     [mp_std, rn_std],
     ['Model Parallel', 'Single GPU'],
     'mp_vs_rn.png')

结果如下所示。模型并行相较于单 GPU 训练的模型，训练时间开销多出 4.02/3.75-1=7% 左右。当然，这存在优化空间，因为多 GPU 中，每一时刻只有一个 GPU 进行训练，其他闲置。而在中间数据转移过程中，又消耗一定的时间。

模型并行 VS 单 GPU

输入流水线

解决上面的问题的最直接的方式就是使用流水线技术，即 GPU-0 输出到 GPU-1 之后，在 GPU-1 训练的同时，GPU-0 接收下一批数据，这样就可以多 GPU 同时执行了。

下面，我们将 120 个样本的 batch 再次细分，分为 20 张样本每份的小 batch。由于 Pytorch 同步启动 CUDA 操作，因此，该操作不需要使用额外的多线程来处理。

class PipelineParallelResNet50(ModelParallelResNet50):
    def __init__(self, split_size=20, *args, **kwargs):
        super(PipelineParallelResNet50, self).__init__(*args, **kwargs)
        self.split_size = split_size

    def forward(self, x):
        splits = iter(x.split(self.split_size, dim=0))
        s_next = next(splits)
        s_prev = self.seq1(s_next).to('cuda:1')
        ret = []

        for s_next in splits:
            # A. s_prev runs on cuda:1
            s_prev = self.seq2(s_prev)
            ret.append(self.fc(s_prev.view(s_prev.size(0), -1)))

            # B. s_next runs on cuda:0, which can run concurrently with A
            s_prev = self.seq1(s_next).to('cuda:1')

        s_prev = self.seq2(s_prev)
        ret.append(self.fc(s_prev.view(s_prev.size(0), -1)))

        return torch.cat(ret)


setup = "model = PipelineParallelResNet50()"
pp_run_times = timeit.repeat(
    stmt, setup, number=1, repeat=num_repeat, globals=globals())
pp_mean, pp_std = np.mean(pp_run_times), np.std(pp_run_times)

plot([mp_mean, rn_mean, pp_mean],
     [mp_std, rn_std, pp_std],
     ['Model Parallel', 'Single GPU', 'Pipelining Model Parallel'],
     'mp_vs_rn_vs_pp.png')

需要注意的是，device-to-device 的 tensor copy 操作是同步的。如果创建多个数据流，则需要保证 copy 操作以合适的同步方式进行。

在完成 tensor 拷贝之前，对 source tensor 进行写入，或者对 target tensor 进行读写，都可能会导致不可预期的行为。上面的实现中，在源和目标设备中，均只使用了默认的 stream，因此无需额外的强化同步操作。

模型并行 VS 单 GPU VS 流水线模型并行

如上图所示，流水线输入确实加速了训练进程，大约 3.75/2.51-1=49%，但距离 100% 的加速相去甚远。由于我们在流水线并行实现中，引入了一个新的参数 split_sizes，但是并不知晓其对训练时间的影响。

直觉上来说，使用一个小的 split_sizes 将会导致许多微小的 CUDA 内核的启动，而使用较大的 split_sizes，则会导致较长的空闲时间。下面是一个搜索最佳 split_sizes 的实验。

means = []
stds = []
split_sizes = [1, 3, 5, 8, 10, 12, 20, 40, 60]

for split_size in split_sizes:
    setup = "model = PipelineParallelResNet50(split_size=%d)" % split_size
    pp_run_times = timeit.repeat(
        stmt, setup, number=1, repeat=num_repeat, globals=globals())
    means.append(np.mean(pp_run_times))
    stds.append(np.std(pp_run_times))

fig, ax = plt.subplots()
ax.plot(split_sizes, means)
ax.errorbar(split_sizes, means, yerr=stds, ecolor='red', fmt='ro')
ax.set_ylabel('ResNet50 Execution Time (Second)')
ax.set_xlabel('Pipeline Split Size')
ax.set_xticks(split_sizes)
ax.yaxis.grid(True)
plt.tight_layout()
plt.savefig("split_size_tradeoff.png")
plt.close(fig)

实验结果如下所示：

流水线输入分割份数

如上图所示，最佳的参数为 12，其将导致 3.75/2.43-1=54% 的加速。但这仍存在加速的可能。例如，所有在 cuda:0 上的操作放在默认的 stream 上。这意味着，在下一个 split 上的计算，不能与上一个 split 的 copy 操作进行重叠。然而，由于 next_split 和 prev_plit 是不同的 tensor，因此这不存在问题。

该实现需要在每个 GPU 上使用多个 stream，并且模型中不同的子网络需要使用不同的 stream 管理策略。

远方到底有多远

Pytorch之GPU程序

GPU