From 48ec36e53faca773c1bfafc6b9a416456f183f6c Mon Sep 17 00:00:00 2001 From: hturki Date: Wed, 1 Feb 2023 17:55:07 -0500 Subject: [PATCH 1/5] 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: From db83068fcd7168d1ac7bdf5bacb7b345bf63cc35 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Wed, 1 Feb 2023 15:15:08 -0800 Subject: [PATCH 2/5] Update optimizer docstrings (#1333) --- nerfstudio/engine/optimizers.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/nerfstudio/engine/optimizers.py b/nerfstudio/engine/optimizers.py index 8281dc6de..6878eccc2 100644 --- a/nerfstudio/engine/optimizers.py +++ b/nerfstudio/engine/optimizers.py @@ -34,9 +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. @@ -54,6 +58,7 @@ class AdamOptimizerConfig(OptimizerConfig): _target: Type = torch.optim.Adam weight_decay: float = 0 + """The weight decay to use.""" @dataclass From 3a416a387bf0a145874826ef32c05ea795c0b290 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Wed, 1 Feb 2023 17:33:06 -0800 Subject: [PATCH 3/5] Dont move viewer outputs to cpu (#1335) --- nerfstudio/viewer/server/viewer_utils.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index 44056537f..1a3fefd51 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/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 From d547279a92216e793e20de959fe352c661be81f2 Mon Sep 17 00:00:00 2001 From: hturki Date: Thu, 2 Feb 2023 12:48:45 -0500 Subject: [PATCH 4/5] 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) From ef9e00e6141d2fde83fc344b3e27ac475a8597cd Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Thu, 2 Feb 2023 10:11:41 -0800 Subject: [PATCH 5/5] Check if curl is installed (#1337) * check if curl is installed * lint --- nerfstudio/utils/install_checks.py | 8 ++++++++ scripts/downloads/download_data.py | 3 +++ 2 files changed, 11 insertions(+) diff --git a/nerfstudio/utils/install_checks.py b/nerfstudio/utils/install_checks.py index 23351a1b2..05e505d4e 100644 --- a/nerfstudio/utils/install_checks.py +++ b/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) diff --git a/scripts/downloads/download_data.py b/scripts/downloads/download_data.py index 12a0a0227..89c77391f 100755 --- a/scripts/downloads/download_data.py +++ b/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":