【简介】
顾名思义,自动混合精度即在使用Pytorch进行网络训练/推理时由程序根据用户输入的option_level的值来判断并自动调整模型训练过程中使用的数值精度以达到降低显存使用,提升训练速度等预期作用。
【例子】
官方Github提供:可运行的ImageNet训练实例
对于GAN网络,官方也提供了一个正在完善中的DCGAN
【简单使用和属性】
Amp允许用户轻松地尝试不同的纯和混合精度模式。通过选择“优化级别”或opt_level
,可以选择常用的默认模式。每opt_level
一个都建立了一组属性,这些属性支配Amp对纯或混合精度训练的实施策略。opt_level
通过将特定属性的值直接传递给amp,可以实现对给定行为的更精确地控制amp.initialize
。这些手动指定的值将覆盖由设置的默认值opt_level
控制。
例子:
# 定义模型和优化器,默认值为FP32精度 model = torch.nn.Linear(D_in, D_out).cuda() optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) # 以opt_level = "01"的方式初始化amp model, optimizer = amp.initialize(model, optimizer, opt_level="O1") ... # 改写loss.backward()如下: with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() ...
使用了Amp后,用户就不应该手动使用.half()缩小模型或者数据的精度。Amp希望用户从现有的默认(FP32)脚本开始,添加与Amp API对应的三行,并开始以混合精度进行训练。也可以禁用Amp,在这种情况下,原始脚本的行为将与以前完全相同。这样,就不会存在Amp API的风险,并且会带来很多潜在的性能优势。
notes:
由于无需手动转换模型(除了call
amp.initialize
)或输入数据,因此遵循新API的脚本可以在不同的opt-levels
之间切换,而无需进行任何其他更改。
函数参数:
-
cast_model_type
:将模型的参数和缓冲区强制转换为所需的类型。 -
patch_torch_functions
:修补所有Torch函数和Tensor方法,以执行Tensor Core友好的操作,例如FP16中的GEMM和卷积,以及任何受益于FP32的FP32精度的操作。 -
keep_batchnorm_fp32
:为了提高精度并启用cudnn batchnorm(可提高性能),即使模型的其余部分是FP16,将batchnorm权重保留在FP32中通常也很有益。 -
master_weights
:保持FP32主权重以与任何FP16型权重一起使用。FP32主权重由优化器步进,以提高精度并捕获小的梯度。 -
loss_scale
:如果loss_scale
是浮动值,则将此值用作静态(固定)损失标度。如果loss_scale
是字符串"dynamic"
,则随时间自适应地调整损耗比例。动态损耗标度调整由Amp自动执行。
通常来说,用户不用手动在函数中去指定这些属性,而是选择一个opt_level
,选择opt_level
后函数将为您设置它们。当然,用户可以选择将kwargs参数手动替代传递(不用opt_level)。
如果您尝试覆盖对所选内容没有意义的参数,则opt_level
Amp会引发错误并给出解释。例如,选择opt_level="O1"
与替代组合参数master_weights=True
是没有意义的。 O1
在Torch函数而非模型权重周围插入强制转换。当数据,激活和权重流经修补的功能时,它们会即时迁移到其他位置。因此,模型权重本身可以(并且应该)保留为FP32,并且无需维护单独的FP32主权重。
【opt_levels】
Amp将opt_levels分为四个常用的组合“00”,“01”,“02”,“03”
-
O0
而O3
不是真正的混合精度,但它们是建立精度和速度的基线。 -
O1
和O2
为具有混合精度不同的实现。在实验中可以两者都尝试,看看哪个选项可以为模型提供最佳的加速和准确性。
00:FP32训练
传入模型精度应该已经是FP32,因此这很可能是空操作(相当于不执行)。 O0
建立准确度基准很有用。
对应传入的参数值为:
cast_model_type = torch.float32 patch_torch_functions = False keep_batchnorm_fp32 = None master_weights = False loss_scale = 1.0
01:混合精度(推荐用于典型用途)
修补所有Torch函数和Tensor方法,以根据白名单-黑名单模型转换其输入。白名单操作(例如,对Tensor Core友好的操作,如GEMM和卷积)在FP16精度下进行。受益于FP32精度(例如softmax)的黑名单操作在FP32中执行。 O1
除非被覆盖,否则还使用动态损耗定标。
对应传入的参数值为:
cast_model_type = None patch_torch_functions = True keep_batchnorm_fp32 = None master_weights = None loss_scale = "dynamic"
02:“大部分都用FP16”混合精度
O2
将模型权重转换为FP16,修补模型的forward
方法以将输入数据转换为FP16,将批范数保留在FP32中,保持FP32主权重,更新优化程序,param_groups
以便optimizer.step()
直接作用于FP32权重(随后为FP32主权重-> FP16模型)权重副本(如有必要),并实现动态损耗定标(除非被覆盖)。不像O1
,O2
不修补Torch函数或Tensor方法。
对应传入的参数值为:
cast_model_type = torch.float16 patch_torch_functions = False keep_batchnorm_fp32 = Ture master_weights = True loss_scale = "dynamic"
03:FP16训练
O3
可能无法实现真正的混合精度选项O1
和的稳定性O2
。但是,为模型建立速度基线是很有用的O1
,O2
可以将和的性能基线进行比较。如果您的模型使用批处理归一化,则可以O3
使用附加的属性替代 keep_batchnorm_fp32=True
(如前所述,启用cudnn batchnorm)来尝试建立“快速训练” 。
对应传入的参数值为:
cast_model_type = torch.float16 patch_torch_functions = False keep_batchnorm_fp32 = False master_weights = False loss_scale = "1.0
【初始化API】
apex.amp.initialize(models, optimizers=None, enabled=True, opt_level='O1', cast_model_type=None, patch_torch_functions=None, keep_batchnorm_fp32=None, master_weights=None, loss_scale=None, cast_model_outputs=None, num_losses=1, verbosity=1, min_loss_scale=None, max_loss_scale=16777216.0)
根据选择opt_level
和覆盖的属性(如果有)初始化模型,优化器以及Torch张量和函数名称空间。
amp.initialize
在完成模型和优化器的构造之后,但在通过任何DistributedDataParallel包装器发送模型之前,应调用。请参阅Imagenet示例中的分布式培训。
当前,它amp.initialize
只能被调用一次,尽管它可以处理任意数量的模型和优化器(请参阅相应的Advanced Amp Usage主题)。如果您认为用例需要amp.initialize
多次调用,请 告诉我们。
任何不是的属性关键字参数都None
将被解释为手动替代。
为避免必须重写脚本中的其他任何内容,请命名返回的模型/优化器以替换传递的模型/优化器,如下面的代码示例中所示。
model, optim = amp.initialize(model, optim,...) model, [optim1, optim2] = amp.initialize(model, [optim1, optim2],...)[model1, model2], optim = amp.initialize([model1, model2], optim,...) [model1, model2], [optim1, optim2] = amp.initialize([model1, model2], [optim1, optim2],...) # This is not an exhaustive list of the cross product of options that are possible, # just a set of examples. model, optim = amp.initialize(model, optim, opt_level="O0") model, optim = amp.initialize(model, optim, opt_level="O0", loss_scale="dynamic"|128.0|"128.0") model, optim = amp.initialize(model, optim, opt_level="O1") # uses "loss_scale="dynamic" default model, optim = amp.initialize(model, optim, opt_level="O1", loss_scale=128.0|"128.0") model, optim = amp.initialize(model, optim, opt_level="O2") # uses "loss_scale="dynamic" default model, optim = amp.initialize(model, optim, opt_level="O2", loss_scale=128.0|"128.0") model, optim = amp.initialize(model, optim, opt_level="O2", keep_batchnorm_fp32=True|False|"True"|"False") model, optim = amp.initialize(model, optim, opt_level="O3") # uses loss_scale=1.0 default model, optim = amp.initialize(model, optim, opt_level="O3", loss_scale="dynamic"|128.0|"128.0") model, optim = amp.initialize(model, optim, opt_level="O3", keep_batchnorm_fp32=True|False|"True"|"False")
参数及其说明
-
model(torch.nn.Module或torch.nn.Modules列表)–要修改/发布的模型。
-
optimizers(可选,torch.optim.Optimizer或torch.optim.Optimizers列表)–修改/发布的优化器。训练所需,推理则可选。
-
enabled(bool,optional ,default = True)–如果为False,则使所有Amp调用均不进行操作,因此您的脚本应像不存在Amp一样运行。
-
opt_level(str,optional ,default =“ O1”)–纯或混合精度优化级别。可接受的值为“ O0”,“ O1”,“ O2”和“ O3”,上面已详细说明。
-
cast_model_type(
torch.dtype
,可选,默认= None)–可选属性覆盖,请参见上文。 -
patch_torch_functions(bool,optional ,default = None)–可选属性覆盖。
-
keep_batchnorm_fp32(bool或str,可选的,默认= None)–可选的属性覆盖。如果作为字符串传递,则必须为字符串“ True”或“ False”。
-
master_weights(bool,optional ,default = None)–可选属性覆盖。
-
loss_scale(float或str,optional ,default = None)–可选属性覆盖。如果作为字符串传递,则必须是代表数字的字符串,例如“ 128.0”或字符串“ dynamic”。
-
cast_model_outputs(torch.dpython:type ,optional ,default = None)–确保模型输出始终转换为特定类型的选项
opt_level
。 -
num_losses(int,optional ,default = 1)–可以提前告知Amp您打算使用多少次损失/向后传递的选项。与
loss_id
to的参数 结合使用时amp.scale_loss
,可使Amp在每次损失/向后传递中使用不同的损失标度,从而可以提高稳定性。有关示例,请参见“高级功放用法”下的“多个型号/优化器/损耗” 。如果num_losses
将其保留为1,则Amp仍将支持多个损失/向后传递,但对所有损失/损失使用单个全局损失等级。 -
verbosity(int,default = 1)–设置为0以抑制与Amp有关的输出。
-
min_loss_scale(float,default = None)–为可以通过动态损失缩放选择的损失缩放值设置下限。默认值是None表示不施加底线。如果不使用动态损耗定标,则忽略min_loss_scale。
-
max_loss_scale(float,默认= 2。** 24)–为可以通过动态损耗缩放选择的损耗比例值设置一个上限。如果不使用动态损耗定标,则忽略max_loss_scale。
返回值
根据修改的模型和优化程序opt_level
。如果models
或optimizers
args是列表,则相应的返回值也将是列表。
【优化损失】
apex.amp.scale_loss(loss, optimizers, loss_id=0, model=None, delay_unscale=False, delay_overflow_check=False)
在上下文管理器入口上,创建。 产生,以便用户可以调用:
scaled_loss = (loss.float())*current loss scale
scaled_loss
scaled_loss.backward()
在上下文管理器出口(如果为delay_unscale=False
)上,将检查渐变的infs / NaNs并进行无比例缩放,以便optimizer.step()
可以调用它。
notes
如果放缩器是使用显式FP32主PARAMS(这是默认的
opt_level=O2
,并且也通过提供手动启用master_weights=True
到amp.initialize
)是未缩放之前的任何FP16梯度被复制到主FP32梯度。optimizer.step()
然后将未缩放的主渐变应用于主参数。
警告
如果Amp使用显式FP32主参数,则仅FP32主渐变不会缩放。
.grad
上下文管理器退出后,任何FP16模型参数的直接属性都将保持缩放。这种微妙之处会影响渐变修剪。有关最佳做法,请参阅“高级放缩器用法”下的“渐变限制”
参数
-
loss –通常为标量张量。在
scaled_loss
该上下文管理的产量仅仅是loss.float()*loss_scale
,所以原则上loss
为你调用可以有一个以上的元素,只要backward()
在scaled_loss
上下文管理主体内适当。 -
optimizers–当前向后传递为其创建渐变的所有优化器。必须是优化程序或从早先调用所返回的优化程序列表
amp.initialize
。例如,与多个优化器一起使用,请参阅“高级放缩器使用情况”下的“多个模型/优化器/损耗” 。 -
loss_id(int,optional ,default = 0)–与
num_losses
参数结合使用时amp.initialize
,可使Amp对每次损失使用不同的损失等级。loss_id
必须为0到0之间的整数,该整数num_losses
告诉Amp当前的反向传递使用了哪种损耗。有关示例,请参见“高级功放用法”下的“多个型号/优化器/损耗” 。如果loss_id
未指定,则Amp将使用默认的全局损耗缩放器进行此向后传递。 -
model(torch.nn.Module ,可选,默认值为None)–当前未使用,保留以用于将来的优化。
-
delay_unscale(bool,optional ,default = False)–
delay_unscale
永远不需要,False
强烈建议使用默认值。如果为True
,则Amp将不会在上下文管理器出口上对渐变进行缩放或执行model-> master渐变副本。delay_unscale=True
是次要的忍者性能优化,可能会导致奇怪的陷阱(尤其是多个模型/优化器/损失),因此只有在知道自己在做什么的情况下才使用它。“高级放缩器用法”下的“迭代之间的梯度累积” 说明了可以(但不必)使用此CAN的情况。
警告
如果
delay_unscale
是True
给定的向后传递,optimizer.step()
则在上下文管理器退出后仍无法调用,并且必须等待另一个稍后向后的上下文管理器调用,并delay_unscale
保留为False。
【apex.amp.master_params】
生成器表达式,它迭代由拥有的参数optimizer
。
参数
优化器–先前从返回的优化器amp.initialize
。
checkpoint(检查点)
为了正确保存和加载您的放大器训练,我们引入了amp.state_dict()
,其中包含所有loss_scaler
以及它们对应的未跳过步骤,以及amp.load_state_dict()
还原这些属性。
为了获得按位精度,我们建议使用以下工作流程:
# Initializationopt_level = 'O1' model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)# Train your model... # Save checkpoint checkpoint = { 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'amp': amp.state_dict() } torch.save(checkpoint, 'amp_checkpoint.pt') ... # Restore model = ... optimizer = ... checkpoint = torch.load('amp_checkpoint.pt') model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level) model.load_state_dict(checkpoint['model']) optimizer.load_state_dict(checkpoint['optimizer']) amp.load_state_dict(checkpoint['amp']) # Continue training...
请注意,我们建议使用相同的恢复模型opt_level
。另请注意,我们建议在load_state_dict
之后调用方法amp.initialize
。
【高级用例】
统一的Amp API支持跨迭代进行梯度累积,每次迭代进行多次反向传递,多个模型/优化器,自定义/用户定义的autograd函数以及自定义数据批处理类。梯度削波和GAN也需要特殊处理,但是这种处理不需要针对不同的opt_level
s进行更改。可以在此处找到更多详细信息: