Pytorch原生AMP支持使用方法(1.6版本)
AMP:Automatic mixed precision,自动混合精度,可以在神经网络推理过程中,针对不同的层,采用不同的数据精度进行计算,从而实现节省显存和加快速度的目的。
在Pytorch 1.5版本及以前,通过NVIDIA出品的插件apex,可以实现amp功能。
从Pytorch 1.6版本以后,Pytorch将amp的功能吸收入官方库,位于torch.cuda.amp
模块下。
本文为针对官方文档主要内容的简要翻译和自己的理解。
1. Introduction
torch.cuda.amp
提供了对混合精度的支持。为实现自动混合精度训练,需要结合使用如下两个模块:
torch.cuda.amp.autocast
:autocast
主要用作上下文管理器或者装饰器,来确定使用混合精度的范围。torch.cuda.amp.GradScalar
:GradScalar
主要用来完成梯度缩放。
2. Typical Mixed Precision Training
一个典型的amp应用示例如下:
# 定义模型和优化器
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
# 在训练最开始定义GradScalar的实例
scaler = GradScaler()
for epoch in epochs:
for input, target in data:
optimizer.zero_grad()
# 利用with语句,在autocast实例的上下文范围内,进行模型的前向推理和loss计算
with autocast():
output = model(input)
loss = loss_fn(output, target)
# 对loss进行缩放,针对缩放后的loss进行反向传播
# (此部分计算在autocast()作用范围以外)
scaler.scale(loss).backward()
# 将梯度值缩放回原尺度后,优化器进行一步优化
scaler.step(optimizer)
# 更新scalar的缩放信息
scaler.update()
3. Working with Unscaled Gradients
待更新
4. Working with Scaled Gradients
待更新
5. Working with Multiple Models, Losses, and Optimizers
如果模型的Loss计算部分输出多个loss,需要对每一个loss值执行scaler.scale
。
如果网络具有多个优化器,对任一个优化器执行scaler.unscale_
,并对每一个优化器执行scaler.step
。
而scaler.update
只在最后执行一次。
应用示例如下:
scaler = torch.cuda.amp.GradScaler()
for epoch in epochs:
for input, target in data:
optimizer0.zero_grad()
optimizer1.zero_grad()
with autocast():
output0 = model0(input)
output1 = model1(input)
loss0 = loss_fn(2 * output0 + 3 * output1, target)
loss1 = loss_fn(3 * output0 - 5 * output1, target)
scaler.scale(loss0).backward(retain_graph=True)
scaler.scale(loss1).backward()
# 选择其中一个优化器执行显式的unscaling
scaler.unscale_(optimizer0)
# 对每一个优化器执行scaler.step
scaler.step(optimizer0)
scaler.step(optimizer1)
# 完成所有梯度更新后,执行一次scaler.update
scaler.update()
6. Working with Multiple GPUs
针对多卡训练的情况,只影响autocast
的使用方法,GradScaler
的用法与之前一致。
6.1 DataParallel in a single process
在每一个不同的cuda设备上,torch.nn.DataParallel
在不同的进程中执行前向推理,而autocast只在当前进程中生效,因此,如下方式的调用是不生效的:
model = MyModel()
dp_model = nn.DataParallel(model)
# 在主进程中设置autocast
with autocast():
# dp_model的内部进程并不会对autocast生效
output = dp_model(input)
# loss的计算在主进程中执行,autocast可以生效,但由于前面执行推理时已经失效,因此整体上是不正确的
loss = loss_fn(output)
有效的调用方式如下所示:
# 方法1:在模型构建中,定义forwar函数时,采用装饰器方式
MyModel(nn.Module):
...
@autocast()
def forward(self, input):
...
# 方法2:在模型构建中,定义forwar函数时,采用上下文管理器方式
MyModel(nn.Module):
...
def forward(self, input):
with autocast():
...
# DataParallel的使用方式不变
model = MyModel().cuda()
dp_model = nn.DataParallel(model)
# 在模型执行推理时,由于前面模型定义时的修改,在各cuda设备上的子进程中autocast生效
# 在执行loss计算是,在主进程中,autocast生效
with autocast():
output = dp_model(input)
loss = loss_fn(output)
6.2 DistributedDataParallel, one GPU per process
torch.nn.parallel.DistributedDataParallel
在官方文档中推荐每个GPU执行一个实例的方法,以达到最好的性能表现。
在这种模式下,DistributedDataParallel
内部并不会再启动子进程,因此对于autocast
和GradScaler
的使用都没有影响,与典型示例保持一致。
6.3 DistributedDataParallel, multiple GPUs per process
与DataParallel
的使用相同,在模型构建时,对forward函数的定义方式进行修改,保证autocast在进程内部生效。