From 7269a7ee1639cdb6d287988a648f53075ca8e1ea Mon Sep 17 00:00:00 2001 From: Sid Mehta Date: Fri, 15 Sep 2023 09:21:09 -0700 Subject: [PATCH] Add Comet Logging (#2431) * initial commit * fixed naming * added comet_ml as depedency * added comet to docs * Update custom_gui.md * Update models.md * Update benchmarking.md * added Comet to ReadMe * added decorator * fixed link and added comet examples * fixed link and added comet example * fixed format * fixed type ignore error * fixed readme typo * fixed typo --- README.md | 2 +- .../debugging_tools/benchmarking.md | 2 +- docs/developer_guides/pipelines/models.md | 2 +- docs/developer_guides/viewer/custom_gui.md | 2 +- docs/quickstart/first_nerf.md | 4 +- nerfstudio/configs/experiment_config.py | 15 ++++--- nerfstudio/engine/trainer.py | 11 ++--- nerfstudio/utils/decorators.py | 2 +- nerfstudio/utils/writer.py | 42 ++++++++++++++++--- pyproject.toml | 1 + 10 files changed, 57 insertions(+), 26 deletions(-) diff --git a/README.md b/README.md index aa5825a23..b6f05f9be 100644 --- a/README.md +++ b/README.md @@ -274,7 +274,7 @@ ns-train nerfacto --help ### Tensorboard / WandB / Viewer -We support three different methods to track training progress, using the viewer, [tensorboard](https://www.tensorflow.org/tensorboard), and [Weights and Biases](https://wandb.ai/site). You can specify which visualizer to use by appending `--vis {viewer, tensorboard, wandb, viewer+wandb, viewer+tensorboard}` to the training command. Simultaneously utilizing the viewer alongside wandb or tensorboard may cause stuttering issues during evaluation steps. The viewer only works for methods that are fast (ie. nerfacto, instant-ngp), for slower methods like NeRF, use the other loggers. +We support four different methods to track training progress, using the viewer[tensorboard](https://www.tensorflow.org/tensorboard), [Weights and Biases](https://wandb.ai/site), and ,[Comet](https://comet.com/?utm_source=nerf&utm_medium=referral&utm_content=github). You can specify which visualizer to use by appending `--vis {viewer, tensorboard, wandb, comet viewer+wandb, viewer+tensorboard, viewer+comet}` to the training command. Simultaneously utilizing the viewer alongside wandb or tensorboard may cause stuttering issues during evaluation steps. The viewer only works for methods that are fast (ie. nerfacto, instant-ngp), for slower methods like NeRF, use the other loggers. # Learn More diff --git a/docs/developer_guides/debugging_tools/benchmarking.md b/docs/developer_guides/debugging_tools/benchmarking.md index b38c09277..bf97db4b5 100644 --- a/docs/developer_guides/debugging_tools/benchmarking.md +++ b/docs/developer_guides/debugging_tools/benchmarking.md @@ -16,7 +16,7 @@ Simply replace the arguments in brackets with the correct arguments. - `-m {METHOD_NAME}`: Name of the method you want to benchmark (e.g. `nerfacto`, `mipnerf`). - `-s`: Launch a single job per GPU. -- `-v {VIS}`: Use another visualization than wandb, which is the default. Only other option is tensorboard. +- `-v {VIS}`: Use another visualization than wandb, which is the default. Other options are comet & tensorboard. - `{GPU_LIST}`: (optional) Specify the list of gpus you want to use on your machine space separated. for instance, if you want to use GPU's 0-3, you will need to pass in `0 1 2 3`. If left empty, the script will automatically find available GPU's and distribute training jobs on the available GPUs. :::{admonition} Tip diff --git a/docs/developer_guides/pipelines/models.md b/docs/developer_guides/pipelines/models.md index 8afa76d3a..4fb1cb3f2 100644 --- a/docs/developer_guides/pipelines/models.md +++ b/docs/developer_guides/pipelines/models.md @@ -55,7 +55,7 @@ class Model: """Process a RayBundle object and return RayOutputs describing quanties for each ray.""" def get_metrics_dict(self, outputs, batch): - """Returns metrics dictionary which will be plotted with wandb or tensorboard.""" + """Returns metrics dictionary which will be plotted with comet, wandb or tensorboard.""" def get_loss_dict(self, outputs, batch, metrics_dict=None): """Returns a dictionary of losses to be summed which will be your loss.""" diff --git a/docs/developer_guides/viewer/custom_gui.md b/docs/developer_guides/viewer/custom_gui.md index e11381d02..2ae41fa6a 100644 --- a/docs/developer_guides/viewer/custom_gui.md +++ b/docs/developer_guides/viewer/custom_gui.md @@ -59,7 +59,7 @@ class MyModel(Model): **Writing to the element** -You can write to a viewer element in Python, which provides a convenient way to track values in your code without the need for wandb/tensorboard or relying on `print` statements. +You can write to a viewer element in Python, which provides a convenient way to track values in your code without the need for comet/wandb/tensorboard or relying on `print` statements. ```python self.custom_value.value = x diff --git a/docs/quickstart/first_nerf.md b/docs/quickstart/first_nerf.md index a88183d9e..137bd0b23 100644 --- a/docs/quickstart/first_nerf.md +++ b/docs/quickstart/first_nerf.md @@ -96,9 +96,9 @@ ns-train nerfacto nerfstudio-data --help Each script will have some other minor quirks (like the training script dataparser subcommand needing to come after the model subcommand), read up on them [here](../reference/cli/index.md). -## Tensorboard / WandB / Viewer +## Comet / Tensorboard / WandB / Viewer -We support three different methods to track training progress, using the viewer, [tensorboard](https://www.tensorflow.org/tensorboard), and [Weights and Biases](https://wandb.ai/site). You can specify which visualizer to use by appending `--vis {viewer, tensorboard, wandb, viewer+wandb, viewer+tensorboard}` to the training command. Simultaneously utilizing the viewer alongside wandb or tensorboard may cause stuttering issues during evaluation steps. The viewer only works for methods that are fast (ie. nerfacto, instant-ngp), for slower methods like NeRF, use the other loggers. +We support four different methods to track training progress, using the viewer [tensorboard](https://www.tensorflow.org/tensorboard), [Weights and Biases](https://wandb.ai/site), and [Comet](https://comet.com/?utm_source=nerf&utm_medium=referral&utm_content=nerf_docs). You can specify which visualizer to use by appending `--vis {viewer, tensorboard, wandb, viewer+wandb, viewer+tensorboard, viewer+comet}` to the training command. Simultaneously utilizing the viewer alongside wandb or tensorboard may cause stuttering issues during evaluation steps. The viewer only works for methods that are fast (ie. nerfacto, instant-ngp), for slower methods like NeRF, use the other loggers. ## Evaluating Runs diff --git a/nerfstudio/configs/experiment_config.py b/nerfstudio/configs/experiment_config.py index 1fa3598f6..0da8893ed 100644 --- a/nerfstudio/configs/experiment_config.py +++ b/nerfstudio/configs/experiment_config.py @@ -22,13 +22,7 @@ from pathlib import Path from typing import Any, Dict, Literal, Optional import yaml - -from nerfstudio.configs.base_config import ( - InstantiateConfig, - LoggingConfig, - MachineConfig, - ViewerConfig, -) +from nerfstudio.configs.base_config import InstantiateConfig, LoggingConfig, MachineConfig, ViewerConfig from nerfstudio.configs.config_utils import to_immutable_dict from nerfstudio.engine.optimizers import OptimizerConfig from nerfstudio.engine.schedulers import SchedulerConfig @@ -68,7 +62,9 @@ class ExperimentConfig(InstantiateConfig): } ) """Dictionary of optimizer groups and their schedulers""" - vis: Literal["viewer", "wandb", "tensorboard", "viewer+wandb", "viewer+tensorboard", "viewer_beta"] = "wandb" + vis: Literal[ + "viewer", "wandb", "tensorboard", "comet", "viewer+wandb", "viewer+tensorboard", "viewer+comet", "viewer_beta" + ] = "wandb" """Which visualizer to use.""" data: Optional[Path] = None """Alias for --pipeline.datamanager.data""" @@ -95,6 +91,9 @@ class ExperimentConfig(InstantiateConfig): """Checks if tensorboard is enabled.""" return ("tensorboard" == self.vis) | ("viewer+tensorboard" == self.vis) + def is_comet_enabled(self) -> bool: + return ("comet" == self.vis) | ("viewer+comet" == self.vis) + def set_timestamp(self) -> None: """Dynamically set the experiment timestamp""" if self.timestamp == "{timestamp}": diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index 94256809c..c8ae50043 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -27,11 +27,6 @@ from threading import Lock from typing import Dict, List, Literal, Optional, Tuple, Type, cast import torch -from rich import box, style -from rich.panel import Panel -from rich.table import Table -from torch.cuda.amp.grad_scaler import GradScaler - from nerfstudio.configs.experiment_config import ExperimentConfig from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes, TrainingCallbackLocation @@ -44,6 +39,10 @@ from nerfstudio.utils.rich_utils import CONSOLE from nerfstudio.utils.writer import EventName, TimeWriter from nerfstudio.viewer.server.viewer_state import ViewerState from nerfstudio.viewer_beta.viewer import Viewer as ViewerBetaState +from rich import box, style +from rich.panel import Panel +from rich.table import Table +from torch.cuda.amp.grad_scaler import GradScaler TRAIN_INTERATION_OUTPUT = Tuple[torch.Tensor, Dict[str, torch.Tensor], Dict[str, torch.Tensor]] TORCH_DEVICE = str @@ -197,6 +196,7 @@ class Trainer: writer.setup_event_writer( self.config.is_wandb_enabled(), self.config.is_tensorboard_enabled(), + self.config.is_comet_enabled(), log_dir=writer_log_path, experiment_name=self.config.experiment_name, project_name=self.config.project_name, @@ -326,6 +326,7 @@ class Trainer: (self.config.is_viewer_enabled() or self.config.is_viewer_beta_enabled()) and not self.config.is_tensorboard_enabled() and not self.config.is_wandb_enabled() + and not self.config.is_comet_enabled() ): string: str = ( "[NOTE] Not running eval iterations since only viewer is enabled.\n" diff --git a/nerfstudio/utils/decorators.py b/nerfstudio/utils/decorators.py index 06f28f60a..d6fd3f417 100644 --- a/nerfstudio/utils/decorators.py +++ b/nerfstudio/utils/decorators.py @@ -66,7 +66,7 @@ def check_eval_enabled(func: Callable) -> Callable: def wrapper(self, *args, **kwargs): ret = None - if self.config.is_wandb_enabled() or self.config.is_tensorboard_enabled(): + if self.config.is_wandb_enabled() or self.config.is_tensorboard_enabled() or self.config.is_comet_enabled(): ret = func(self, *args, **kwargs) return ret diff --git a/nerfstudio/utils/writer.py b/nerfstudio/utils/writer.py index 8d3b07b1a..59e6c5a71 100644 --- a/nerfstudio/utils/writer.py +++ b/nerfstudio/utils/writer.py @@ -24,16 +24,16 @@ from pathlib import Path from time import time from typing import Any, Dict, List, Optional, Union +import comet_ml import torch import wandb from jaxtyping import Float -from torch import Tensor -from torch.utils.tensorboard import SummaryWriter - from nerfstudio.configs import base_config as cfg from nerfstudio.utils.decorators import check_main_thread, decorate_all from nerfstudio.utils.printing import human_format from nerfstudio.utils.rich_utils import CONSOLE +from torch import Tensor +from torch.utils.tensorboard import SummaryWriter def to8b(x): @@ -150,7 +150,7 @@ def put_time(name: str, duration: float, step: int, avg_over_steps: bool = True, put_scalar(name, duration, step) if update_eta: - ## NOTE: eta should be called with avg train iteration time + # NOTE: eta should be called with avg train iteration time remain_iter = GLOBAL_BUFFER["max_iter"] - step remain_time = remain_iter * GLOBAL_BUFFER["events"][name]["avg"] put_scalar(EventName.ETA, remain_time, step) @@ -185,7 +185,7 @@ def setup_local_writer(config: cfg.LoggingConfig, max_iter: int, banner_messages else: CONSOLE.log("disabled local writer") - ## configure all the global buffer basic information + # configure all the global buffer basic information GLOBAL_BUFFER["max_iter"] = max_iter GLOBAL_BUFFER["max_buffer_size"] = config.max_buffer_size GLOBAL_BUFFER["steps_per_log"] = config.steps_per_log @@ -203,6 +203,7 @@ def is_initialized(): def setup_event_writer( is_wandb_enabled: bool, is_tensorboard_enabled: bool, + is_comet_enabled: bool, log_dir: Path, experiment_name: str, project_name: str = "nerfstudio-project", @@ -214,6 +215,11 @@ def setup_event_writer( banner_messages: list of messages to always display at bottom of screen """ using_event_writer = False + + if is_comet_enabled: + curr_writer = CometWriter(log_dir=log_dir, experiment_name=experiment_name, project_name=project_name) + EVENT_WRITERS.append(curr_writer) + using_event_writer = True if is_wandb_enabled: curr_writer = WandbWriter(log_dir=log_dir, experiment_name=experiment_name, project_name=project_name) EVENT_WRITERS.append(curr_writer) @@ -225,7 +231,7 @@ def setup_event_writer( if using_event_writer: string = f"logging events to: {log_dir}" else: - string = "Disabled tensorboard/wandb event writers" + string = "Disabled comet/tensorboard/wandb event writers" CONSOLE.print(f"[bold yellow]{string}") @@ -346,6 +352,30 @@ class TensorboardWriter(Writer): self.tb_writer.add_text("config", str(config_dict)) +@decorate_all([check_main_thread]) +class CometWriter(Writer): + """Comet_ML Writer Class""" + + def __init__(self, log_dir: Path, experiment_name: str, project_name: str = "nerfstudio-project"): + self.experiment = comet_ml.Experiment(project_name=project_name) + if experiment_name != "unnamed": + self.experiment.set_name(experiment_name) + + def write_image(self, name: str, image: Float[Tensor, "H W C"], step: int) -> None: + self.experiment.log_image(image, name, step=step) + + def write_scalar(self, name: str, scalar: Union[float, torch.Tensor], step: int) -> None: + self.experiment.log_metric(name, scalar, step) + + def write_config(self, name: str, config_dict: Dict[str, Any], step: int): + """Function that writes out the config to Comet + + Args: + config: config dictionary to write out + """ + self.experiment.log_parameters(config_dict, step=step) + + def _cursorup(x: int): """utility tool to move the cursor up on the terminal diff --git a/pyproject.toml b/pyproject.toml index 1347f954b..add21cf6d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,6 +16,7 @@ classifiers = [ dependencies = [ "appdirs>=1.4", "av>=9.2.0", + "comet_ml>=3.33.8", "cryptography>=38", "tyro>=0.5.3", "gdown>=4.6.0",