pytorch半精度,混合精度,单精度训练的区别amp.initialize

mixed_precision = True
try:  # Mixed precision training https://github.com/NVIDIA/apex
    from apex import amp
except:
    mixed_precision = False  # not installed

 model, optimizer = amp.initialize(model, optimizer, opt_level='O1', verbosity=1)

为了帮助提高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。

默认训练方式是 单精度float32

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()

半精度 model(img.half())

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()

接下来是混合精度的实现,这里主要用到Apex的amp工具。代码修改为:
加上这一句封装,model, optimizer = amp.initialize(model, optimizer, opt_level=“O1”)

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进行回传。

需要注意以下几点:

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

在使用混合精度进行计算时,最关键的参数是opt_level。他一共含有四种设置值:‘00’,‘01’,‘02’,‘03’。实际上整个amp.initialize的输入参数很多:

但是在实际使用过程中发现,设置opt_level即可,这也是文档中例子的使用方法,甚至在不同的opt_level设置条件下,其他的参数会变成无效。(已知BUG:使用‘01’时设置keep_batchnorm_fp32的值会报错)

概括起来:00相当于原始的单精度训练。01在大部分计算时采用半精度,但是所有的模型参数依然保持单精度,对于少数单精度较好的计算(如softmax)依然保持单精度。02相比于01,将模型参数也变为半精度。03基本等于最开始实验的全半精度的运算。值得一提的是,不论在优化过程中,模型是否采用半精度,保存下来的模型均为单精度模型,能够保证模型在其他应用中的正常使用。这也是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)

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

你可能感兴趣的:(pytorch半精度,混合精度,单精度训练的区别amp.initialize)