From a8bce87ada9124bf5f2e0b73d5d918f8cf9f1c35 Mon Sep 17 00:00:00 2001
From: Terrance Wang <terrance.3512@gmail.com>
Date: Thu, 25 May 2023 09:27:58 -0700
Subject: [PATCH] pass grad scaler from trainer to pipeline (#1996)

* pass grad sclaer from trainer to pipeline

* lint

* add grad scaler to arg doc

---------

Co-authored-by: Matthew Tancik <tancik@berkeley.edu>
---
 nerfstudio/engine/trainer.py          | 6 +++++-
 nerfstudio/pipelines/base_pipeline.py | 5 +++++
 nerfstudio/pipelines/dynamic_batch.py | 4 +++-
 3 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py
index 4a3a8565c..e786b9ab2 100644
--- a/nerfstudio/engine/trainer.py
+++ b/nerfstudio/engine/trainer.py
@@ -148,7 +148,11 @@ class Trainer:
                 'inference': does not load any dataset into memory
         """
         self.pipeline = self.config.pipeline.setup(
-            device=self.device, test_mode=test_mode, world_size=self.world_size, local_rank=self.local_rank
+            device=self.device,
+            test_mode=test_mode,
+            world_size=self.world_size,
+            local_rank=self.local_rank,
+            grad_scaler=self.grad_scaler,
         )
         self.optimizers = self.setup_optimizers()
 
diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py
index ebc1c2498..1e1e0fab5 100644
--- a/nerfstudio/pipelines/base_pipeline.py
+++ b/nerfstudio/pipelines/base_pipeline.py
@@ -35,6 +35,7 @@ from rich.progress import (
 from torch import nn
 from torch.nn import Parameter
 from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.cuda.amp.grad_scaler import GradScaler
 
 from nerfstudio.configs import base_config as cfg
 from nerfstudio.data.datamanagers.base_datamanager import (
@@ -222,6 +223,7 @@ class VanillaPipeline(Pipeline):
             'inference': does not load any dataset into memory
         world_size: total number of machines available
         local_rank: rank of current machine
+        grad_scaler: gradient scaler used in the trainer
 
     Attributes:
         datamanager: The data manager that will be used
@@ -235,6 +237,7 @@ class VanillaPipeline(Pipeline):
         test_mode: Literal["test", "val", "inference"] = "val",
         world_size: int = 1,
         local_rank: int = 0,
+        grad_scaler: Optional[GradScaler] = None,
     ):
         super().__init__()
         self.config = config
@@ -250,6 +253,8 @@ class VanillaPipeline(Pipeline):
             scene_box=self.datamanager.train_dataset.scene_box,
             num_train_data=len(self.datamanager.train_dataset),
             metadata=self.datamanager.train_dataset.metadata,
+            device=device,
+            grad_scaler=grad_scaler,
         )
         self.model.to(device)
 
diff --git a/nerfstudio/pipelines/dynamic_batch.py b/nerfstudio/pipelines/dynamic_batch.py
index fda064c07..65679d001 100644
--- a/nerfstudio/pipelines/dynamic_batch.py
+++ b/nerfstudio/pipelines/dynamic_batch.py
@@ -17,9 +17,10 @@ A pipeline that dynamically chooses the number of rays to sample.
 """
 
 from dataclasses import dataclass, field
-from typing import Literal, Type
+from typing import Literal, Type, Optional
 
 import torch
+from torch.cuda.amp.grad_scaler import GradScaler
 
 from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager
 from nerfstudio.pipelines.base_pipeline import VanillaPipeline, VanillaPipelineConfig
@@ -50,6 +51,7 @@ class DynamicBatchPipeline(VanillaPipeline):
         test_mode: Literal["test", "val", "inference"] = "val",
         world_size: int = 1,
         local_rank: int = 0,
+        grad_scaler: Optional[GradScaler] = None,
     ):
         super().__init__(config, device, test_mode, world_size, local_rank)
         assert isinstance(