def enable_gradient_clipping(model: Model,
grad_clipping: Optional[float]) -> None:
if grad_clipping is not None:
for parameter in model.parameters():
if parameter.requires_grad:
parameter.register_hook(lambda grad: nn_util.clamp_tensor(
grad, minimum=-grad_clipping, maximum=grad_clipping))
def rescale_gradients(model: Model,
grad_norm: Optional[float] = None) -> Optional[float]:
"""
Performs gradient rescaling. Is a no-op if gradient rescaling is not enabled.
"""
if grad_norm:
parameters_to_clip = [
p for p in model.parameters() if p.grad is not None
]
return sparse_clip_norm(parameters_to_clip, grad_norm)
return None