FP16
稍微介绍一下，FP16，FP32，BF16。
FP32是单精度浮点数，8 bit表示指数，23bit表示小数。FP16采用5bit表示指数，10bit表示小数。BF采用8bit表示整数，7bit表示小数。所以总结就是，BF16的整数范围等于FP32，但是精度差。FP16的表示范围和精度都低于FP32。
在mmdetction这种框架中，如果要使用FP16，其实只需要一行代码就可以了。
fp16 = dict(loss_scale=512.)
当然，你要使用fp16，首先你的GPU要支持才可以。
接下来这段代码告诉我们，其实fp16_cfg这个东西，决定的是optimizer_config。
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
# 如果我们设置了，则会生成一个Fp16OptimizerHook的实例
optimizer_config = Fp16OptimizerHook(
**cfg.optimizer_config, **fp16_cfg, distributed=False)
else:
# 如果我们没有设置，则正常从config里面读取optimizer_config
# 如设置grad_clip: optimizer_config = dict(grad_clip=dict(max_norm=35, norm_type=2))
optimizer_config = cfg.optimizer_config
# 然后注册训练的hooks,optimizer_config会被当参数传进去
runner.register_training_hooks(cfg.lr_config, optimizer_config,
cfg.checkpoint_config, cfg.log_config)
还可以看看registe_training_hooks这个函数，register_optimizer_hook这个函数。
def register_training_hooks(self,lr_config, optimizer_config=None,
checkpoint_config=None,log_config=None):
self.register_lr_hook(lr_config)
# 这里注册传进来的optimizer_config，其他hook不需要关注
self.register_optimizer_hook(optimizer_config)
self.register_checkpoint_hook(checkpoint_config)
self.register_hook(IterTimerHook())
self.register_logger_hooks(log_config)
def register_optimizer_hook(self, optimizer_config):
if optimizer_config is None:
return
# 如果是dict，则生成OptimizerHook的实例，正常的反传和更新参数
if isinstance(optimizer_config, dict):
optimizer_config.setdefault('type', 'OptimizerHook')
hook = mmcv.build_from_cfg(optimizer_config, HOOKS)
# 如果不是dict，那就是我们之前传进来的Fp16OptimizerHook的实例了
else:
hook = optimizer_config
# 注册这个hook，就是添加到hook_list里,待训练的时候某个指定时间节点使用
self.register_hook(hook)
在Fp16OptimizerHook的实现上，需要注意的三个事情是：
1）需要拷贝一份FP32权重用来更新，在FP16这个表示下，梯度和权重都是基于半精度来表示和存储的。那么在运算的时候，很有可能运算结果就小到FP16的极限表示能力以下了。所以这里要采用fp32来进行运算。所以用fp32来进行step操作。
2）需要将loss放大，这也是那个scale的作用。如果梯度比较小的话，FP16的表示能力就会把梯度变成0。
3）torch的权重存储在model中，可以通过parameters()来获取。optimizer为了更新权重，所以在param_groups里面也存了一份（这里共享了内存）。model里的FP16，optimizer里面的FP32数据类型都不一样了。所以这里要解耦开，用FP16在model里存，但是用FP32在optimizer里面进行更新。这里的说法是，权重的内存和特征图比起来，其实没有那么多。特征图都是FP16的，所以不用担心会造成很多额外的存储上的overhead。
class Fp16OptimizerHook(OptimizerHook):
def __init__(self,
grad_clip,
coalesce=True,
bucket_size_mb=-1,
loss_scale=512,
distributed=True
):
self.grad_clip = grad_clip
self.coalesce = coalesce
self.bucket_size = bucket_size
self.loss_scale = loss_scale
self.distributed = distributed
def before_run(self, runner):
# param_groups是torch，optimizer的成员变量
# dict，keys有‘params’，‘learning_rate’,'momentum', 'weight_decay'等信息
# 本来是与model等权重同一块内存，但是现在重新开了一块出来，这就是解耦
runner.optimizer.param_groups = copy.deepcopy(
runner.optimizer.param_groups
)
这个函数主要是把model等weigths存储空间削成一半
wrap_fp16_model(runner_model)
def after_train_iter(self,runner):
#解除耦合之后，要做两次梯度清零
runner.model.zero_grad()
runner.optimizer.zeor_grad()
# 在backward之前，乘上一个系数，还是在避免超出最小表示范围。
scaled_loss = runner.outputs['loss'] * self.loss_scale
scaled_loss.backward()
fp32_weigts=[]
for param_group in runner.optimizer_groups:
fp32_weigts += param_group['param']
# copy FP16的梯度值进FP32的梯度值里面。
self.copy_grads_to_fp32(runner.model, fp32_weights)
# 针对分布式训练
if self.distributed:
allreduce_grads(fp32_weights, self.coalesce, self.bucket_size_mb)
for param in fp32_weights:
if param.grad is not None:
param.grad.div_(self.loss_scale)
if self.grad_clip is not None:
self.clip_grads(fp32_weights)
# optimizer更新参数，利用FP32进行计算
runner.optimizer.step()
# 算完之后将optimizer的数值拷贝到model里面，以FP16进行存储
self.copy_params_to_fp16(runner.model, fp32.weights)
def copy_grads_to_fp32(self, fp16_net, fp32_weights):
"""Copy gradients from fp16 model to fp32 weight copy."""
for fp32_param, fp16_param in zip(fp32_weights, fp16_net.parameters()):
if fp16_param.grad is not None:
if fp32_param.grad is None:
fp32_param.grad = fp32_param.data.new(fp32_param.size())
fp32_param.grad.copy_(fp16_param.grad)
# 这里直接copy就好了
def copy_params_to_fp16(self, fp16_net, fp32_weights):
"""Copy updated params from fp32 weight copy to fp16 model."""
for fp16_param, fp32_param in zip(fp16_net.parameters(), fp32_weights):
fp16_param.data.copy_(fp32_param.data)
