稍微介紹一下，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)