From d547279a92216e793e20de959fe352c661be81f2 Mon Sep 17 00:00:00 2001 From: hturki Date: Thu, 2 Feb 2023 12:48:45 -0500 Subject: [PATCH] add option to log training gradients (#1334) --- nerfstudio/engine/trainer.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index f24ea5ccd..2d4a45395 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -77,6 +77,8 @@ class TrainerConfig(ExperimentConfig): load_step: Optional[int] = None """Optionally specify model step to load from; if none, will find most recent model in load_dir.""" load_config: Optional[Path] = None + """Optionally log gradients during training""" + log_gradients: bool = False class Trainer: @@ -188,7 +190,6 @@ class Trainer: step = 0 for step in range(self._start_step, self._start_step + num_iterations): with TimeWriter(writer, EventName.ITER_TRAIN_TIME, step=step) as train_t: - self.pipeline.train() # training callbacks before the training iteration @@ -370,6 +371,18 @@ class Trainer: loss = functools.reduce(torch.add, loss_dict.values()) self.grad_scaler.scale(loss).backward() # type: ignore self.optimizers.optimizer_scaler_step_all(self.grad_scaler) + + if self.config.log_gradients: + total_grad = 0 + for tag, value in self.pipeline.model.named_parameters(): + assert tag != "Total" + if value.grad is not None: + grad = value.grad.norm() + metrics_dict[f"Gradients/{tag}"] = grad + total_grad += grad + + metrics_dict["Gradients/Total"] = total_grad + self.grad_scaler.update() self.optimizers.scheduler_step_all(step)