def _track_grad_norm(self, trainer: "pl.Trainer") -> None: if trainer.track_grad_norm == -1: return grad_norm_dict = grad_norm(trainer.lightning_module, trainer.track_grad_norm, trainer.logger.group_separator) if grad_norm_dict: prev_fx = trainer.lightning_module._current_fx_name trainer.lightning_module._current_fx_name = "on_before_optimizer_step" trainer.lightning_module.log_grad_norm(grad_norm_dict) trainer.lightning_module._current_fx_name = prev_fx
def _track_grad_norm(self, trainer: "pl.Trainer") -> None: if trainer.track_grad_norm == -1: return kwargs = {} if len(trainer.loggers) == 1: kwargs["group_separator"] = trainer.loggers[0].group_separator grad_norm_dict = grad_norm(trainer.lightning_module, trainer.track_grad_norm, **kwargs) if grad_norm_dict: prev_fx = trainer.lightning_module._current_fx_name trainer.lightning_module._current_fx_name = "on_before_optimizer_step" trainer.lightning_module.log_grad_norm(grad_norm_dict) trainer.lightning_module._current_fx_name = prev_fx
def test_grad_norm(norm_type, expected): """Test utility function for computing the p-norm of individual parameter groups and norm in total.""" class Model(nn.Module): def __init__(self): super().__init__() self.param0 = nn.Parameter(torch.rand(3)) self.param1 = nn.Parameter(torch.rand(2, 1)) self.param0.grad = torch.tensor([-1.0, 2.0, -3.0]) self.param1.grad = torch.tensor([[-4.0], [5.0]]) # param without grad should not contribute to norm self.param2 = nn.Parameter(torch.rand(1)) model = Model() norms = grad_norm(model, norm_type) expected = {k: round(v, 4) for k, v in expected.items()} assert norms == expected
def _track_and_norm_grad(self, optimizer: torch.optim.Optimizer) -> Dict[str, Tensor]: """Tracks gradient norms and clips the gradients of all parameters optimized by the current optimizer. Args: optimizer: the current optimizer """ # track gradient norms grad_norm_dict = {} can_log = (self.trainer.global_step + 1) % self.trainer.log_every_n_steps == 0 should_track = float(self.trainer.track_grad_norm) > 0 if should_track and can_log: grad_norm_dict = grad_norm(self.trainer.lightning_module, self.trainer.track_grad_norm) # clip gradients self.trainer.accelerator.clip_gradients( optimizer, self.trainer.gradient_clip_val, gradient_clip_algorithm=self.trainer.gradient_clip_algorithm ) return grad_norm_dict
def test_grad_norm_invalid_norm_type(norm_type): with pytest.raises(ValueError, match="`norm_type` must be a positive number or 'inf'"): grad_norm(Mock(), norm_type)