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

@@ -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
 

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

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."""

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

docs/quickstart/first_nerf.md

@@ -96,9 +96,9 @@ ns-train nerfacto <nerfacto optional args> 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
 

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}":

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"

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
 

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
 

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",