From 48ec36e53faca773c1bfafc6b9a416456f183f6c Mon Sep 17 00:00:00 2001
From: hturki <hturki@cs.cmu.edu>
Date: Wed, 1 Feb 2023 17:55:07 -0500
Subject: [PATCH] Add support for gradient clipping (#1331)

implement optional gradient clipping
---
 nerfstudio/engine/optimizers.py | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/nerfstudio/engine/optimizers.py b/nerfstudio/engine/optimizers.py
index c97b10580..8281dc6de 100644
--- a/nerfstudio/engine/optimizers.py
+++ b/nerfstudio/engine/optimizers.py
@@ -18,7 +18,7 @@ Optimizers class.
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any, Dict, List, Type
+from typing import Any, Dict, List, Optional, Type
 
 import torch
 from torch.cuda.amp.grad_scaler import GradScaler
@@ -36,6 +36,7 @@ class OptimizerConfig(base_config.PrintableConfig):
     _target: Type = torch.optim.Adam
     lr: float = 0.0005
     eps: float = 1e-08
+    max_norm: Optional[float] = None
 
     # TODO: somehow make this more generic. i dont like the idea of overriding the setup function
     # but also not sure how to go about passing things into predefined torch objects.
@@ -43,6 +44,7 @@ class OptimizerConfig(base_config.PrintableConfig):
         """Returns the instantiated object using the config."""
         kwargs = vars(self).copy()
         kwargs.pop("_target")
+        kwargs.pop("max_norm")
         return self._target(params, **kwargs)
 
 
@@ -73,9 +75,11 @@ class Optimizers:
         self.config = config
         self.optimizers = {}
         self.schedulers = {}
+        self.parameters = {}
         for param_group_name, params in param_groups.items():
             lr_init = config[param_group_name]["optimizer"].lr
             self.optimizers[param_group_name] = config[param_group_name]["optimizer"].setup(params=params)
+            self.parameters[param_group_name] = params
             if config[param_group_name]["scheduler"]:
                 self.schedulers[param_group_name] = config[param_group_name]["scheduler"].setup(
                     optimizer=self.optimizers[param_group_name], lr_init=lr_init
@@ -109,13 +113,20 @@ class Optimizers:
         Args:
             grad_scaler: GradScaler to use
         """
-        for _, optimizer in self.optimizers.items():
+        for param_group, optimizer in self.optimizers.items():
+            max_norm = self.config[param_group]["optimizer"].max_norm
+            if max_norm is not None:
+                grad_scaler.unscale_(optimizer)
+                torch.nn.utils.clip_grad_norm_(self.parameters[param_group], max_norm)
             grad_scaler.step(optimizer)
 
     def optimizer_step_all(self):
         """Run step for all optimizers."""
-        for _, optimizer in self.optimizers.items():
+        for param_group, optimizer in self.optimizers.items():
             # note that they key is the parameter name
+            max_norm = self.config[param_group]["optimizer"].max_norm
+            if max_norm is not None:
+                torch.nn.utils.clip_grad_norm_(self.parameters[param_group], max_norm)
             optimizer.step()
 
     def scheduler_step_all(self, step: int) -> None: