Apex混合精度加速

栏目: Python · 发布时间: 5年前

内容简介:为了帮助提高Pytorch的训练效率,英伟达提供了混合精度训练工具Apex。号称能够在不降低性能的情况下,将模型训练的速度提升2-4倍,训练显存消耗减少为之前的一半。该项目开源于:https://github.com/NVIDIA/apex ,文档地址是:https://nvidia.github.io/apex/index.html该工具提供了三个功能,amp、parallel和normalization。由于目前该工具还是0.1版本,功能还是很基础的,在最后一个normalization功能中只提供了L

为了帮助提高Pytorch的训练效率,英伟达提供了混合精度训练工具Apex。号称能够在不降低性能的情况下,将模型训练的速度提升2-4倍,训练显存消耗减少为之前的一半。该项目开源于:https://github.com/NVIDIA/apex ,文档地址是:https://nvidia.github.io/apex/index.html

工具 提供了三个功能,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:
	out = model(img)
	loss = LOSS(out, label)
	loss.backward()
	optimizer.step()
	optimizer.zero_grad()

这种状态下,我测试的模型的单卡显存占用为: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:
	out = model(img.half())
	loss = LOSS(out, label)
	loss.backward()
	optimizer.step()
	optimizer.zero_grad()

即将模型和输入的数据转为半精度即可。这样将达到最高的速度和最小的模型体积。采用该方法,我的测试模型单卡显存占用为:4978MB,速度为34s/40iter

对比这两种情况,显存占用的确是大大降低了,之前模型训练需要的显存位8466*4 = 33864MB,采用半精度训练后显存占用降低为:4978*4 = 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:
	out = model(img)
	loss = LOSS(out, label)
	# loss.backward()
	with amp.scale_loss(loss, optimizer) as scaled_loss:
    	scaled_loss.backward()

	optimizer.step()
	optimizer.zero_grad()

实际流程为:调用amp.initialize按照预定的opt_level对model和optimizer进行设置。在计算loss时使用amp.scale_loss进行回传。

需要注意以下几点:

  1. 在调用amp.initialize之前,模型需要放在GPU上,也就是需要调用cuda()或者to()。
  2. 在调用amp.initialize之前,模型不能调用任何分布式设置函数。
  3. 此时输入数据不需要在转换为半精度。

在使用混合精度进行计算时,最关键的参数是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的值会报错)

下表展示了不同设置的差异:

设置编号 00 01 02 03
cast_model_type torch.float32 None torch.float16 torch.float16
patch_torch_functions False True False False
keep_batchnorm_fp32 None None True False
master_weights False None True False
loss_scale 1.0 “dynamic” “dynamic” 1.0

概括起来: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):
	classname = m.__class__.__name__
    if classname.find('BatchNorm') != -1:
    	m.eval().half()

model.apply(fix_bn)

实际测试下来,最后的模型准确度上感觉差别不大,可能有轻微下降;时间上变化不大,这可能会因不同的模型有差别;显存开销上确实有很大的降低。


以上所述就是小编给大家介绍的《Apex混合精度加速》,希望对大家有所帮助,如果大家有任何疑问请给我留言,小编会及时回复大家的。在此也非常感谢大家对 码农网 的支持!

查看所有标签

猜你喜欢:

本站部分资源来源于网络,本站转载出于传递更多信息之目的,版权归原作者或者来源机构所有,如转载稿涉及版权问题,请联系我们

从零到百亿

从零到百亿

〔美〕卡罗·白朗、Karel Baloun、卡罗·白朗 / 译言网 / 中国书籍出版社 / 2007 / 15.00元

Facebook是一个发源于哈佛大学、为全美大学生服务的社交网站。按照流量,这个网站在世界范围排名第8名;按照价值,业界对Facebook公司的估值超过10亿美元。Facebook创建于2004年2月,这样的高速增长成为当今互联网发展的一个奇迹。 《Inside Facebook》这本书是原作者卡罗·白朗(Karel Baloun)作为Facebook的第一位高级软件开发人员之一,在Face......一起来看看 《从零到百亿》 这本书的介绍吧!

UNIX 时间戳转换
UNIX 时间戳转换

UNIX 时间戳转换

RGB HSV 转换
RGB HSV 转换

RGB HSV 互转工具

HEX CMYK 转换工具
HEX CMYK 转换工具

HEX CMYK 互转工具