Allow for kwargs to be passed to trackers (#542)

* Allow for kwarg passing to trackers

Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
Co-authored-by: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com>

Author: 3 people

src/accelerate/accelerator.py

@@ -927,7 +927,7 @@ def wait_for_everyone(self):
         """
         wait_for_everyone()
 
-    def init_trackers(self, project_name: str, config: Optional[dict] = None):
+    def init_trackers(self, project_name: str, config: Optional[dict] = None, init_kwargs: Optional[dict] = {}):
         """
         Initializes a run for all trackers stored in `self.log_with`, potentially with starting configurations
 
@@ -936,6 +936,12 @@ def init_trackers(self, project_name: str, config: Optional[dict] = None):
                 The name of the project. All trackers will save their data based on this
             config (`dict`, *optional*):
                 Optional starting configuration to be logged.
+            init_kwargs (`dict`, *optional*):
+                A nested dictionary of kwargs to be passed to a specific tracker's `__init__` function. Should be
+                formatted like this:
+                ```python
+                {"wandb": {"tags": ["tag_a", "tag_b"]}}
+                ```
         """
         self.trackers = []
         for tracker in self.log_with:
@@ -946,14 +952,16 @@ def init_trackers(self, project_name: str, config: Optional[dict] = None):
                 tracker_init = LOGGER_TYPE_TO_CLASS[str(tracker)]
                 if getattr(tracker_init, "requires_logging_directory"):
                     # We can skip this check since it was done in `__init__`
-                    self.trackers.append(tracker_init(project_name, self.logging_dir))
+                    self.trackers.append(
+                        tracker_init(project_name, self.logging_dir, **init_kwargs.get(str(tracker), {}))
+                    )
                 else:
-                    self.trackers.append(tracker_init(project_name))
+                    self.trackers.append(tracker_init(project_name, **init_kwargs.get(str(tracker), {})))
         if config is not None:
             for tracker in self.trackers:
                 tracker.store_init_configuration(config)
 
-    def log(self, values: dict, step: Optional[int] = None):
+    def log(self, values: dict, step: Optional[int] = None, log_kwargs: Optional[dict] = {}):
         """
         Logs `values` to all stored trackers in `self.trackers`.
 
@@ -962,10 +970,16 @@ def log(self, values: dict, step: Optional[int] = None):
                 Values should be a dictionary-like object containing only types `int`, `float`, or `str`.
             step (`int`, *optional*):
                 The run step. If included, the log will be affiliated with this step.
+            log_kwargs (`dict`, *optional*):
+                A nested dictionary of kwargs to be passed to a specific tracker's `log` function. Should be formatted
+                like this:
+                ```python
+                {"wandb": {"tags": ["tag_a", "tag_b"]}}
+                ```
         """
         if self.is_main_process:
             for tracker in self.trackers:
-                tracker.log(values, step=step)
+                tracker.log(values, step=step, **log_kwargs.get(tracker.name, {}))
 
     def end_training(self):
         """

src/accelerate/tracking.py

@@ -52,8 +52,16 @@ def get_available_trackers():
 class GeneralTracker(object, metaclass=ABCMeta):
     """
     A base Tracker class to be used for all logging integration implementations.
+
+    Each function should take in `**kwargs` that will automatically be passed in from a base dictionary provided to
+    [`Accelerator`]
     """
 
+    @abstractproperty
+    def name(self):
+        "String representation of the python class name"
+        pass
+
     @abstractproperty
     def requires_logging_directory(self):
         """
@@ -75,7 +83,7 @@ def store_init_configuration(self, values: dict):
         pass
 
     @abstractmethod
-    def log(self, values: dict, step: Optional[int]):
+    def log(self, values: dict, step: Optional[int], **kwargs):
         """
         Logs `values` to the current run. Base `log` implementations of a tracking API should go in here, along with
         special behavior for the `step parameter.
@@ -105,14 +113,17 @@ class TensorBoardTracker(GeneralTracker):
             The name of the experiment run
         logging_dir (`str`, `os.PathLike`):
             Location for TensorBoard logs to be stored.
+        kwargs:
+            Additional key word arguments passed along to the `tensorboard.SummaryWriter.__init__` method.
     """
 
+    name = "tensorboard"
     requires_logging_directory = True
 
-    def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]]):
+    def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]], **kwargs):
         self.run_name = run_name
         self.logging_dir = os.path.join(logging_dir, run_name)
-        self.writer = tensorboard.SummaryWriter(self.logging_dir)
+        self.writer = tensorboard.SummaryWriter(self.logging_dir, **kwargs)
         logger.info(f"Initialized TensorBoard project {self.run_name} logging to {self.logging_dir}")
         logger.info(
             "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
@@ -131,7 +142,7 @@ def store_init_configuration(self, values: dict):
         self.writer.flush()
         logger.info("Stored initial configuration hyperparameters to TensorBoard")
 
-    def log(self, values: dict, step: Optional[int] = None):
+    def log(self, values: dict, step: Optional[int] = None, **kwargs):
         """
         Logs `values` to the current run.
 
@@ -141,14 +152,17 @@ def log(self, values: dict, step: Optional[int] = None):
                 `str` to `float`/`int`.
             step (`int`, *optional*):
                 The run step. If included, the log will be affiliated with this step.
+            kwargs:
+                Additional key word arguments passed along to either `SummaryWriter.add_scaler`,
+                `SummaryWriter.add_text`, or `SummaryWriter.add_scalers` method based on the contents of `values`.
         """
         for k, v in values.items():
             if isinstance(v, (int, float)):
-                self.writer.add_scalar(k, v, global_step=step)
+                self.writer.add_scalar(k, v, global_step=step, **kwargs)
             elif isinstance(v, str):
-                self.writer.add_text(k, v, global_step=step)
+                self.writer.add_text(k, v, global_step=step, **kwargs)
             elif isinstance(v, dict):
-                self.writer.add_scalars(k, v, global_step=step)
+                self.writer.add_scalars(k, v, global_step=step, **kwargs)
         self.writer.flush()
         logger.info("Successfully logged to TensorBoard")
 
@@ -167,13 +181,16 @@ class WandBTracker(GeneralTracker):
     Args:
         run_name (`str`):
             The name of the experiment run.
+        kwargs:
+            Additional key word arguments passed along to the `wandb.init` method.
     """
 
+    name = "wandb"
     requires_logging_directory = False
 
-    def __init__(self, run_name: str):
+    def __init__(self, run_name: str, **kwargs):
         self.run_name = run_name
-        self.run = wandb.init(project=self.run_name)
+        self.run = wandb.init(project=self.run_name, **kwargs)
         logger.info(f"Initialized WandB project {self.run_name}")
         logger.info(
             "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
@@ -191,7 +208,7 @@ def store_init_configuration(self, values: dict):
         wandb.config.update(values)
         logger.info("Stored initial configuration hyperparameters to WandB")
 
-    def log(self, values: dict, step: Optional[int] = None):
+    def log(self, values: dict, step: Optional[int] = None, **kwargs):
         """
         Logs `values` to the current run.
 
@@ -201,8 +218,10 @@ def log(self, values: dict, step: Optional[int] = None):
                 `str` to `float`/`int`.
             step (`int`, *optional*):
                 The run step. If included, the log will be affiliated with this step.
+            kwargs:
+                Additional key word arguments passed along to the `wandb.log` method.
         """
-        self.run.log(values, step=step)
+        self.run.log(values, step=step, **kwargs)
         logger.info("Successfully logged to WandB")
 
     def finish(self):
@@ -222,13 +241,16 @@ class CometMLTracker(GeneralTracker):
     Args:
         run_name (`str`):
             The name of the experiment run.
+        kwargs:
+            Additional key word arguments passed along to the `Experiment.__init__` method.
     """
 
+    name = "comet_ml"
     requires_logging_directory = False
 
-    def __init__(self, run_name: str):
+    def __init__(self, run_name: str, **kwargs):
         self.run_name = run_name
-        self.writer = Experiment(project_name=run_name)
+        self.writer = Experiment(project_name=run_name, **kwargs)
         logger.info(f"Initialized CometML project {self.run_name}")
         logger.info(
             "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
@@ -246,7 +268,7 @@ def store_init_configuration(self, values: dict):
         self.writer.log_parameters(values)
         logger.info("Stored initial configuration hyperparameters to CometML")
 
-    def log(self, values: dict, step: Optional[int] = None):
+    def log(self, values: dict, step: Optional[int] = None, **kwargs):
         """
         Logs `values` to the current run.
 
@@ -256,16 +278,19 @@ def log(self, values: dict, step: Optional[int] = None):
                 `str` to `float`/`int`.
             step (`int`, *optional*):
                 The run step. If included, the log will be affiliated with this step.
+            kwargs:
+                Additional key word arguments passed along to either `Experiment.log_metric`, `Experiment.log_other`,
+                or `Experiment.log_metrics` method based on the contents of `values`.
         """
         if step is not None:
             self.writer.set_step(step)
         for k, v in values.items():
             if isinstance(v, (int, float)):
-                self.writer.log_metric(k, v, step=step)
+                self.writer.log_metric(k, v, step=step, **kwargs)
             elif isinstance(v, str):
-                self.writer.log_other(k, v)
+                self.writer.log_other(k, v, **kwargs)
             elif isinstance(v, dict):
-                self.writer.log_metrics(v, step=step)
+                self.writer.log_metrics(v, step=step, **kwargs)
         logger.info("Successfully logged to CometML")
 
     def finish(self):

tests/test_tracking.py

@@ -100,7 +100,8 @@ def test_init_trackers(self):
         project_name = "test_project_with_config"
         accelerator = Accelerator(log_with="wandb")
         config = {"num_iterations": 12, "learning_rate": 1e-2, "some_boolean": False, "some_string": "some_value"}
-        accelerator.init_trackers(project_name, config)
+        kwargs = {"wandb": {"tags": ["my_tag"]}}
+        accelerator.init_trackers(project_name, config, kwargs)
         accelerator.end_training()
         # The latest offline log is stored at wandb/latest-run/*.wandb
         for child in Path(f"{self.tmpdir}/wandb/latest-run").glob("*"):
@@ -116,6 +117,7 @@ def test_init_trackers(self):
         self.assertEqual(self.get_value_from_log("learning_rate", cleaned_log), "0.01")
         self.assertEqual(self.get_value_from_log("some_boolean", cleaned_log), "false")
         self.assertEqual(self.get_value_from_log("some_string", cleaned_log), "some_value")
+        self.assertIn("my_tag", cleaned_log)
 
     def test_log(self):
         project_name = "test_project_with_log"
@@ -214,6 +216,7 @@ class MyCustomTracker(GeneralTracker):
         "some_string",
     ]
 
+    name = "my_custom_tracker"
     requires_logging_directory = False
 
     def __init__(self, dir: str):