| 该 工具 提供了三个功能，amp、parallel和normalization。由于目前该工具还是0.1版本，功能还是很基础的，在最后一个normalization功能中只提供了LayerNorm层的复现，实际上在后续的使用过程中会发现，出现问题最多的是pytorch的BN层。 |
|
| 第二个工具是pytorch的分布式训练的复现，在文档中描述的是和pytorch中的实现等价，在代码中可以选择任意一个使用，实际使用过程中发现，在使用混合精度训练时，使用Apex复现的parallel工具，能避免一些bug。 |
|
| 以上提到的两个工具和pytorch中对应的在使用方法上可以等价替换，这里就不再过多的介绍了。个人实验中，Layer Norm层使用的是Pytorch的版本，分布式训练使用的是Apex的版本。 |
|
| 接下来是最关键的混合精度工具amp的介绍。首先介绍一下什么是混合精度： |
|
| Pytorch模型中默认使用的是32位的float值进行数学运算，所有模型也保存为float32类型，故一般将以float32数据类型进行训练的过程称为单精度计算。对应的以float16为数据类型进行训练的过程称为半精度计算。相比于单精度运算，半精度训练速度更快，模型大小更小但是有些计算需要更高精度的计算,另外单精度的运算稳定性更好。因此我们可以通过混合精度的计算来减少训练开销。 |
|
| 首先一个默认的原始训练流程为： |
|
| import torch |
| model = torch.nn.Linear(D_in, D_out) |
| optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) |
| for img, label in dataloader: |
|
|
|
|
|
| 这种状态下，我测试的模型的单卡显存占用为：8466MB，速度为18s/40iter 如果想使用半精度的训练，可以直接将流程换为： |
|
| import torch |
| model = torch.nn.Linear(D_in, D_out).half() |
| optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) |
| for img, label in dataloader: |
|
|
|
|
|
| 即将模型和输入的数据转为半精度即可。这样将达到最高的速度和最小的模型体积。采用该方法，我的测试模型单卡显存占用为：4978MB，速度为34s/40iter |
|
| 对比这两种情况，显存占用的确是大大降低了，之前模型训练需要的显存位84664 = 33864MB，采用半精度训练后显存占用降低为：49784 = 11948MB。至于速度的降低暂时还没找到原因。根据GitHub的讨论区中的讨论，有一个初步的判断是：CUDNN的加速只支持float32，将模型转为半精度之后，不能再使用cudnn进行加速，这对模型的计算速度影响较大。 |
|
| 接下来是混合精度的实现，这里主要用到Apex的amp工具。代码修改为： |
|
| import torch |
| model = torch.nn.Linear(D_in, D_out).cuda() |
| optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) |
| model, optimizer = amp.initialize(model, optimizer, opt_level="O1") |
|
| for img, label in dataloader: |
|
|
|
|
|
|
|
|
| 实际流程为：调用amp.initialize按照预定的opt_level对model和optimizer进行设置。在计算loss时使用amp.scale_loss进行回传。 |
|
| 需要注意以下几点： |
|
| 在调用amp.initialize之前，模型需要放在GPU上，也就是需要调用cuda()或者to()。 |
| 在调用amp.initialize之前，模型不能调用任何分布式设置函数。 |
| 此时输入数据不需要在转换为半精度。 |
| 在使用混合精度进行计算时，最关键的参数是opt_level。他一共含有四种设置值：‘00’，‘01’，‘02’，‘03’。实际上整个amp.initialize的输入参数很多： |
|
| def initialize( |
| models, |
| optimizers=None, |
| enabled=True, |
| opt_level=None, |
| 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, |
| ): |
| """ |
| Args: |
| models (torch.nn.Module or list of torch.nn.Modules): Models to modify/cast. |
| optimizers (optional, torch.optim.Optimizer or list of torch.optim.Optimizers): Optimizers to modify/cast. |
| REQUIRED for training, optional for inference. |
| enabled (bool, optional, default=True): If False, renders all Amp calls no-ops, so your script |
| should run as if Amp were not present. |
| opt_level (str, required): Pure or mixed precision optimization level. Accepted values are |
| "O0", "O1", "O2", and "O3", explained in detail above. |
cast_model_type (torch.dtype, optional, default=None): Optional property override, see |
| above. |
| patch_torch_functions (bool, optional, default=None): Optional property override. |
| keep_batchnorm_fp32 (bool or str, optional, default=None): Optional property override. If |
| passed as a string, must be the string "True" or "False". |
| master_weights (bool, optional, default=None): Optional property override. |
| loss_scale (float or str, optional, default=None): Optional property override. If passed as a string, |
| must be a string representing a number, e.g., "128.0", or the string "dynamic". |
| cast_model_outputs (torch.dtype, optional, default=None): Option to ensure that the outputs |
of your model(s) are always cast to a particular type regardless of opt_level. |
| num_losses (int, optional, default=1): Option to tell Amp in advance how many losses/backward |
passes you plan to use. When used in conjunction with the loss_id argument to |
amp.scale_loss, enables Amp to use a different loss scale per loss/backward pass, |
| which can improve stability. See "Multiple models/optimizers/losses" |
under Advanced Amp Usage_ for examples. If num_losses is left to 1, Amp will still |
| support multiple losses/backward passes, but use a single global loss scale |
| for all of them. |
| verbosity (int, default=1): Set to 0 to suppress Amp-related output. |
| 但是在实际使用过程中发现，设置opt_level即可，这也是文档中例子的使用方法，甚至在不同的opt_level设置条件下，其他的参数会变成无效。（已知BUG：使用‘01’时设置keep_batchnorm_fp32的值会报错） |
|
| 下表展示了不同设置的差异： |
|
| 设置编号 |
| cast_model_type |
| patch_torch_functions |
| keep_batchnorm_fp32 |
| master_weights |
| loss_scale |
| 概括起来：00相当于原始的单精度训练。01在大部分计算时采用半精度，但是所有的模型参数依然保持单精度，对于少数单精度较好的计算（如softmax）依然保持单精度。02相比于01，将模型参数也变为半精度。03基本等于最开始实验的全半精度的运算。值得一提的是，不论在优化过程中，模型是否采用半精度，保存下来的模型均为单精度模型，能够保证模型在其他应用中的正常使用。这也是Apex的一大卖点。 |
|
| 实际对比使用中01和02设置单卡分别使用的显存量为01:5402MiB，02:5426MiB，基本不分上下。时间和半精度持平，为34s/40iter。 |
|
| 此外，在03设置中，能够允许keep_batchnorm_fp32为True，这样在计算BN层时能够使用CUDNN进行加速，按文档的说法能够达到最高的训练速度。实际测试和设为False没有差别。 |
|
| 既然提到BN层，BN层的优化是目前Apex工具最大的问题所在。下面详细说明一下： |
|
| 在Pytorch中，BN层分为train和eval两种操作。实现时若为单精度网络，会调用CUDNN进行计算加速。常规训练过程中BN层会被设为train。Apex优化了这种情况，通过设置keep_batchnorm_fp32参数，能够保证此时BN层使用CUDNN进行计算，达到最好的计算速度。但是在一些fine tunning场景下，BN层会被设为eval（我的模型就是这种情况）。此时keep_batchnorm_fp32的设置并不起作用，训练会产生数据类型不正确的bug。此时需要人为的将所有BN层设置为半精度，这样将不能使用CUDNN加速。一个设置的参考代码如下： |
|
| def fix_bn(m): |
|
| if classname.find('BatchNorm') != -1: |
|
|
| model.apply(fix_bn) |
| 实际测试下来，最后的模型准确度上感觉差别不大，可能有轻微下降；时间上变化不大，这可能会因不同的模型有差别；显存开销上确实有很大的降低。 |
|
|
|
| 本文来源：码农网 |
| 本文链接：https://www.codercto.com/a/79904.html |
