Merge remote-tracking branch 'origin/main' into generative

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
@@ -34,8 +34,13 @@ class OptimizerConfig(base_config.PrintableConfig):
     """Basic optimizer config with RAdam"""
 
     _target: Type = torch.optim.Adam
+    """The optimizer class to use."""
     lr: float = 0.0005
+    """The learning rate to use."""
     eps: float = 1e-08
+    """The epsilon value to use."""
+    max_norm: Optional[float] = None
+    """The max norm to use for gradient clipping."""
 
     # 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 +48,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)
 
 
@@ -52,6 +58,7 @@ class AdamOptimizerConfig(OptimizerConfig):
 
     _target: Type = torch.optim.Adam
     weight_decay: float = 0
+    """The weight decay to use."""
 
 
 @dataclass
@@ -73,9 +80,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 +118,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:

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:
@@ -189,7 +191,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
@@ -371,6 +372,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)
 

nerfstudio/utils/install_checks.py

@@ -39,3 +39,11 @@ def check_colmap_installed():
         CONSOLE.print("[bold red]Could not find COLMAP. Please install COLMAP.")
         print("See https://colmap.github.io/install.html for installation instructions.")
         sys.exit(1)
+
+
+def check_curl_installed():
+    """Checks if curl is installed."""
+    curl_path = shutil.which("curl")
+    if curl_path is None:
+        CONSOLE.print("[bold red]Could not find [yellow]curl[red], Please install [yellow]curl")
+        sys.exit(1)

nerfstudio/viewer/server/viewer_utils.py

@@ -48,7 +48,6 @@ from nerfstudio.models.base_model import Model
 from nerfstudio.utils import colormaps, profiler, writer
 from nerfstudio.utils.decorators import check_main_thread, decorate_all
 from nerfstudio.utils.io import load_from_json, write_to_json
-from nerfstudio.utils.misc import get_dict_to_torch
 from nerfstudio.utils.writer import GLOBAL_BUFFER, EventName, TimeWriter
 from nerfstudio.viewer.server.subprocess import run_viewer_bridge_server_as_subprocess
 from nerfstudio.viewer.server.utils import (
@@ -163,7 +162,6 @@ class RenderThread(threading.Thread):
             self.exc = e
 
         if outputs:
-            outputs = get_dict_to_torch(outputs)
             self.vis_outputs = outputs
 
         self.state.check_done_render = True

scripts/downloads/download_data.py

@@ -15,6 +15,7 @@ from rich.console import Console
 from typing_extensions import Annotated
 
 from nerfstudio.configs.base_config import PrintableConfig
+from nerfstudio.utils import install_checks
 from nerfstudio.utils.scripts import run_command
 
 CONSOLE = Console(width=120)
@@ -212,6 +213,7 @@ class DNerfDownload(DatasetDownload):
         """Download the D-NeRF dataset (https://github.com/albertpumarola/D-NeRF)."""
         # TODO: give this code the same structure as download_nerfstudio
 
+        install_checks.check_curl_installed()
         final_path = save_dir / Path("dnerf")
         if os.path.exists(final_path):
             shutil.rmtree(str(final_path))
@@ -253,6 +255,7 @@ class PhototourismDownload(DatasetDownload):
     def download(self, save_dir: Path):
         """Download a PhotoTourism dataset: https://www.cs.ubc.ca/~kmyi/imw2020/data.html"""
 
+        install_checks.check_curl_installed()
         if self.capture_name == "all":
             for capture_name in phototourism_downloads:
                 if capture_name != "all":