def maybe_add_gradient_clipping(
cfg: CfgNode, optimizer: torch.optim.Optimizer
) -> torch.optim.Optimizer:
"""
If gradient clipping is enabled through config options, wraps the existing
optimizer instance of some type OptimizerType to become an instance
of the new dynamically created class OptimizerTypeWithGradientClip
that inherits OptimizerType and overrides the `step` method to
include gradient clipping.
Args:
cfg: CfgNode
configuration options
optimizer: torch.optim.Optimizer
existing optimizer instance
Return:
optimizer: torch.optim.Optimizer
either the unmodified optimizer instance (if gradient clipping is
disabled), or the same instance with adjusted __class__ to override
the `step` method and include gradient clipping
"""
if not cfg.SOLVER.CLIP_GRADIENTS.ENABLED:
return optimizer
grad_clipper = _create_gradient_clipper(cfg.SOLVER.CLIP_GRADIENTS)
OptimizerWithGradientClip = _generate_optimizer_class_with_gradient_clipping(
type(optimizer), grad_clipper
)
optimizer.__class__ = OptimizerWithGradientClip
return optimizer
def maybe_add_gradient_clipping(
cfg: CfgNode,
optimizer: torch.optim.Optimizer) -> torch.optim.Optimizer:
if not cfg.SOLVER.CLIP_GRADIENTS.ENABLED:
return optimizer
grad_clipper = _create_gradient_clipper(cfg.SOLVER.CLIP_GRADIENTS)
OptimizerWithGradientClip = _generate_optimizer_class_with_gradient_clipping(
type(optimizer), grad_clipper)
optimizer.__class__ = OptimizerWithGradientClip
return optimizer
