mirror of
https://gitcode.com/gh_mirrors/eas/EasyFace.git
synced 2026-07-20 11:37:47 +00:00
78 lines
3.0 KiB
Python
78 lines
3.0 KiB
Python
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
import logging
|
|
|
|
from modelscope.metainfo import Hooks
|
|
from modelscope.trainers.hooks.builder import HOOKS
|
|
|
|
from .base import OptimizerHook
|
|
|
|
|
|
@HOOKS.register_module(module_name=Hooks.ApexAMPOptimizerHook)
|
|
class ApexAMPOptimizerHook(OptimizerHook):
|
|
"""
|
|
Fp16 optimizer, if torch version is less than 1.6.0,
|
|
you must install apex (https://www.github.com/nvidia/apex) else use torch.cuda.amp by default
|
|
|
|
Args:
|
|
cumulative_iters (int): interval of gradients accumulation. Default: 1
|
|
grad_clip (dict): Default None. Containing keys:
|
|
max_norm (float or int): max norm of the gradients
|
|
norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for infinity norm.
|
|
More details please refer to `torch.nn.utils.clip_grad.clip_grad_norm_`
|
|
loss_keys (str | list): keys list of loss
|
|
opt_level (str): "O0" and "O3" are not true mixed precision,
|
|
but they are useful for establishing accuracy and speed baselines, respectively.
|
|
"O1" and "O2" are different implementations of mixed precision.
|
|
Try both, and see what gives the best speedup and accuracy for your model.
|
|
"""
|
|
def __init__(self,
|
|
cumulative_iters=1,
|
|
grad_clip=None,
|
|
loss_keys='loss',
|
|
opt_level='O1'):
|
|
|
|
super(ApexAMPOptimizerHook, self).__init__(grad_clip=grad_clip,
|
|
loss_keys=loss_keys)
|
|
self.cumulative_iters = cumulative_iters
|
|
self.opt_level = opt_level
|
|
|
|
try:
|
|
from apex import amp
|
|
except ImportError:
|
|
raise ValueError(
|
|
'apex not installed, please install apex from https://www.github.com/nvidia/apex.'
|
|
)
|
|
|
|
def before_run(self, trainer):
|
|
from apex import amp
|
|
|
|
logging.info('open fp16')
|
|
# TODO: fix it should initialze amp with model not wrapper by DDP or DP
|
|
if hasattr(trainer.model, 'module'):
|
|
trainer.model, trainer.optimizer = amp.initialize(
|
|
trainer.model.module,
|
|
trainer.optimizer,
|
|
opt_level=self.opt_level)
|
|
else:
|
|
trainer.model, trainer.optimizer = amp.initialize(
|
|
trainer.model, trainer.optimizer, opt_level=self.opt_level)
|
|
|
|
trainer.optimizer.zero_grad()
|
|
|
|
def after_train_iter(self, trainer):
|
|
for k in self.loss_keys:
|
|
trainer.train_outputs[k] /= self.cumulative_iters
|
|
|
|
from apex import amp
|
|
for k in self.loss_keys:
|
|
with amp.scale_loss(trainer.train_outputs[k],
|
|
trainer.optimizer) as scaled_loss:
|
|
scaled_loss.backward()
|
|
|
|
if self.every_n_iters(trainer, self.cumulative_iters):
|
|
if self.grad_clip is not None:
|
|
self.clip_grads(trainer.model.parameters(), **self.grad_clip)
|
|
|
|
trainer.optimizer.step()
|
|
trainer.optimizer.zero_grad()
|