113 lines
3.8 KiB
Python
113 lines
3.8 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
import logging
|
|
|
|
from torch.nn.utils import clip_grad
|
|
|
|
from modelscope.metainfo import Hooks
|
|
from modelscope.outputs import OutputKeys
|
|
from modelscope.trainers.hooks.builder import HOOKS
|
|
from modelscope.trainers.hooks.hook import Hook
|
|
from modelscope.trainers.hooks.priority import Priority
|
|
|
|
|
|
class OptimizerProcessor:
|
|
|
|
def initialize_optimizer(self, trainer):
|
|
"""Initialize the optimizer.
|
|
|
|
This is a strategic function which can be registered by other hook's function.
|
|
"""
|
|
trainer.optimizer.zero_grad()
|
|
|
|
def before_forward(self, trainer):
|
|
pass
|
|
|
|
def backward(self, trainer, loss_keys, cumulative_iters, grad_clip):
|
|
"""Do module backward, optimizer's step and zero_grad and clip the grads.
|
|
|
|
This is a strategic function which can be registered by other hook's function.
|
|
|
|
Args:
|
|
trainer(`EpochBasedTrainer`): The trainer instance.
|
|
loss_keys(`list`): The list of loss keys.
|
|
cumulative_iters(`int`): The cumulative iters for gradients.
|
|
grad_clip(`dict`): The grad clipping options.
|
|
"""
|
|
for k in loss_keys:
|
|
trainer.train_outputs[k] /= cumulative_iters
|
|
trainer.train_outputs[k].backward()
|
|
|
|
if Hook.every_n_iters(trainer, cumulative_iters):
|
|
if grad_clip is not None:
|
|
self.clip_grads(trainer.model.parameters(), **grad_clip)
|
|
|
|
trainer.optimizer.step()
|
|
trainer.optimizer.zero_grad()
|
|
|
|
@staticmethod
|
|
def clip_grads(params, **clip_args):
|
|
params = list(
|
|
filter(lambda p: p.requires_grad and p.grad is not None, params))
|
|
if len(params) > 0:
|
|
return clip_grad.clip_grad_norm_(params, **clip_args)
|
|
|
|
|
|
@HOOKS.register_module(module_name=Hooks.OptimizerHook)
|
|
class OptimizerHook(Hook):
|
|
"""Optimizer hook
|
|
|
|
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
|
|
"""
|
|
|
|
PRIORITY = Priority.ABOVE_NORMAL
|
|
|
|
def __init__(self,
|
|
cumulative_iters=1,
|
|
grad_clip=None,
|
|
loss_keys=OutputKeys.LOSS,
|
|
**kwargs) -> None:
|
|
if isinstance(loss_keys, str):
|
|
loss_keys = [loss_keys]
|
|
assert isinstance(loss_keys, (tuple, list))
|
|
self.loss_keys = loss_keys
|
|
self.cumulative_iters = cumulative_iters
|
|
self.grad_clip = grad_clip
|
|
self.processor = OptimizerProcessor()
|
|
|
|
def set_processor(self, processor):
|
|
self.processor = processor
|
|
|
|
def before_run(self, trainer):
|
|
trainer.cumulative_iters = self.cumulative_iters
|
|
self.processor.initialize_optimizer(trainer)
|
|
|
|
def before_train_iter(self, trainer):
|
|
self.processor.before_forward(trainer)
|
|
|
|
def after_train_iter(self, trainer):
|
|
self.processor.backward(trainer, self.loss_keys, self.cumulative_iters,
|
|
self.grad_clip)
|
|
|
|
|
|
@HOOKS.register_module(module_name=Hooks.NoneOptimizerHook)
|
|
class NoneOptimizerHook(OptimizerHook):
|
|
|
|
def __init__(self, cumulative_iters=1, grad_clip=None, loss_keys='loss'):
|
|
|
|
super(NoneOptimizerHook, self).__init__(
|
|
grad_clip=grad_clip, loss_keys=loss_keys)
|
|
self.cumulative_iters = cumulative_iters
|
|
|
|
def before_run(self, trainer):
|
|
return
|
|
|
|
def after_train_iter(self, trainer):
|
|
return
|