From c55f5d36b92cb61010de72271bba9a33d2c8c71c Mon Sep 17 00:00:00 2001 From: ksnzh Date: Sun, 18 Jun 2023 22:17:35 +0800 Subject: [PATCH] add gradient accumulation support (#2049) * add gradient accumulation support * fix 2 blank lines * fix possibly unbound variable * update gradient accumulation step with for loop * add seert and ignore pyright check --------- Co-authored-by: Zhang Jian --- nerfstudio/engine/trainer.py | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index 4afb3a69e..0ce05db0b 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -33,8 +33,8 @@ from rich.table import Table from torch.cuda.amp.grad_scaler import GradScaler from nerfstudio.configs.experiment_config import ExperimentConfig -from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes, TrainingCallbackLocation from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager +from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes, TrainingCallbackLocation from nerfstudio.engine.optimizers import Optimizers from nerfstudio.pipelines.base_pipeline import VanillaPipeline from nerfstudio.utils import profiler, writer @@ -82,6 +82,8 @@ class TrainerConfig(ExperimentConfig): """Path to checkpoint file.""" log_gradients: bool = False """Optionally log gradients during training""" + gradient_accumulation_steps: int = 1 + """Number of steps to accumulate gradients over.""" class Trainer: @@ -118,6 +120,7 @@ class Trainer: self.mixed_precision: bool = self.config.mixed_precision self.use_grad_scaler: bool = self.mixed_precision or self.config.use_grad_scaler self.training_state: Literal["training", "paused", "completed"] = "training" + self.gradient_accumulation_steps: int = self.config.gradient_accumulation_steps if self.device == "cpu": self.mixed_precision = False @@ -457,10 +460,15 @@ class Trainer: self.optimizers.zero_grad_all() cpu_or_cuda_str: str = self.device.split(":")[0] - with torch.autocast(device_type=cpu_or_cuda_str, enabled=self.mixed_precision): - _, loss_dict, metrics_dict = self.pipeline.get_train_loss_dict(step=step) - loss = functools.reduce(torch.add, loss_dict.values()) - self.grad_scaler.scale(loss).backward() # type: ignore + assert ( + self.gradient_accumulation_steps > 0 + ), f"gradient_accumulation_steps must be > 0, not {self.gradient_accumulation_steps}" + for _ in range(self.gradient_accumulation_steps): + with torch.autocast(device_type=cpu_or_cuda_str, enabled=self.mixed_precision): + _, loss_dict, metrics_dict = self.pipeline.get_train_loss_dict(step=step) + loss = functools.reduce(torch.add, loss_dict.values()) + loss /= self.gradient_accumulation_steps + self.grad_scaler.scale(loss).backward() # type: ignore self.optimizers.optimizer_scaler_step_all(self.grad_scaler) if self.config.log_gradients: @@ -469,10 +477,10 @@ class Trainer: assert tag != "Total" if value.grad is not None: grad = value.grad.norm() - metrics_dict[f"Gradients/{tag}"] = grad + metrics_dict[f"Gradients/{tag}"] = grad # type: ignore total_grad += grad - metrics_dict["Gradients/Total"] = cast(torch.Tensor, total_grad) + metrics_dict["Gradients/Total"] = cast(torch.Tensor, total_grad) # type: ignore scale = self.grad_scaler.get_scale() self.grad_scaler.update() @@ -481,7 +489,7 @@ class Trainer: self.optimizers.scheduler_step_all(step) # Merging loss and metrics dict into a single output. - return loss, loss_dict, metrics_dict + return loss, loss_dict, metrics_dict # type: ignore @check_eval_enabled @profiler.time_function