FP16

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)

reference

1. https://zhuanlan.zhihu.com/p/114438961