未经授权,严禁转载!
另外,近一年没用过这个了,所以内容自己也记不得了,各位读者朋友们,抱歉不能回答你们的问题。
个人主页:
- 会飞的咸鱼tankzhoufirst.github.io参考:
Optional : Data Parallelismpytorch.org DataParallel layers (multi-GPU, distributed)pytorch.org Model Parallel Best Practicespytorch.org PyTorch 大批量数据在单个或多个 GPU 训练指南www.pytorchtutorial.com (原)PyTorch中使用指定的GPU - darkknightzh - 博客园www.cnblogs.comCUDA 基本使用
查看 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') # 默认为当前设备
还可以通过设备类型加上编号,来创建 device
对象:
device = torch.device('cuda', 0)
device = torch.device('cpu', 0)
配置 CUDA 访问限制
可以通过如下方式,设置当前 Python
脚本可见的 GPU
。
在终端设置
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 代码中设置
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0, 2"
使用函数 set_device
import torch
torch.cuda.set_device(id)
官方建议使用CUDA_VISIBLE_DEVICES
,不建议使用set_device
函数。
单机单卡训练
用 GPU 训练
默认情况下,使用 CPU
训练模型。可以通过如下方式,通过 GPU
进行训练。使用 GPU 时,模型和输入必须位于同一张 GPU 上。
.to(device)
和 .cuda()
的区别如下:
.to()
中的参数必不可少- 对于
module
而言,.to()
是inplace
的,而.cuda()
不是;而对于tensor
而言,两者一致。
注:实测,两者时间消耗持平。
方式 1 :
device = torch.device("cuda:1") # 指定模型训练所在 GPU
# 将 GPU 转移至 GPU
if torch.cuda.is_available() and use_gpu:
net = net.cuda(device) # 默认在第一块 GPU 上训练
# 同时将数据转移至 GPU
if torch.cuda.is_available() and use_gpu:
inputs = inputs.cuda(device)
labels = labels.cuda(device)
方法 2 :
device = torch.device("cuda:1") # 指定模型训练所在 GPU
# 将 GPU 转移至 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
分别训练模型的不同部分,使用同一批数据。
两者对比如下图所示:
数据并行
Pytorch API
【Class 原型】
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]
【返回值】
要进行并行的模型。
【基本使用方式】
>>> 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
- 在每次迭代时,执行如下操作:
- 将
controller
模型复制(broadcast
)到每一个指定的GPU
上 - 将总输入的数据
batch
,进行均分,分别作为各对应副本的输入 (scatter
) - 每个副本独立进行前向传播,并进行反向传播,但只是求取梯度
- 将各副本的梯度汇总(
gather
)到controller
设备,并进行求和 (reduced add
)During the backwards pass, gradients from each replica are summed into the original module
. - 更具总体度,更新
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
进行训练,其他闲置。而在中间数据转移过程中,又消耗一定的时间。
输入流水线
解决上面的问题的最直接的方式就是使用流水线技术,即 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
,因此无需额外的强化同步操作。
如上图所示,流水线输入确实加速了训练进程,大约 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
管理策略。