Repository: huggingface/accelerate
Pull request: Add logging capabilities
Start of an implementation for a convinence API when logging with Accelerate. I've opted to use abstract functions here in the base class so that if people wish to integrate with their own API's and have it be compatible, an error can be raised if they missed an implementation. Let me know what you think so far @sgugger

Closes https://github.com/huggingface/accelerate/issues/291

Problem statement:
Add logging capability
# Feature Request: Logging Capability

## Feature Summary
- Add a method called `Accelerator.log` that takes in dictionary-like values and saves to configured logger
- `Accelerator.__init__` should get a new param of `log_with` that takes a list of strings as potential options or `"all"`
- The default should be `None` and no loggers are configured. If `"all"` is selected, find all installed loggers (similar to `Trainer`)

## Timeline
Supported integrations:
- [x] TensorBoard/X
- [x] WandB
- [x] CometML
- [ ] NeptuneAI
﻿

Hints from issue discussion:
Is it possible to add support for custom loggers too? Like the open-source experiment tracker Aim (others above other than tboard are closed source): https://github.com/aimhubio/aim . 
@ashutoshsaboo essentially this will just be a small convenience wrapper to a few back-ends. For any and all loggers you can simply have them log below the following if statement and it should run without issue:
```python
if accelerator.is_main_process:
  # Do my logging of everything here
```
That's fine, I'm already doing that currently but that's way too much boilerplate code currently. Would I still be able to use Accelerator.log somehow? I was thinking more in the sense of having an interface (maybe something similar to Trainer's callback interface) which Accelerate could also do, and can be extended easily? This helps various lib owners provide their own custom impls that can be used out of the box, rather than everyone having to figure out specifics of any other expt tracker and implement it for themselves. @muellerzr

Commit message: Initial implementation for review

Patch to review:
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 2b623db340d..ac684620d71 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -17,10 +17,12 @@
 import sys
 import warnings
 from contextlib import contextmanager
+from pathlib import Path
 from typing import List, Optional, Union
 
 import torch
 
+from accelerate.accelerate.src.accelerate.tracking import TensorBoardTracker
 from packaging import version
 
 from .checkpointing import load_accelerator_state, load_custom_state, save_accelerator_state, save_custom_state
@@ -30,12 +32,14 @@
 from .state import AcceleratorState, DistributedType, is_deepspeed_available
 from .utils import (
     DeepSpeedPlugin,
+    LoggerType,
     PrecisionType,
     RNGType,
     convert_outputs_to_fp32,
     extract_model_from_parallel,
     gather,
     get_pretty_name,
+    is_tensorboard_available,
     pad_across_processes,
     save,
     wait_for_everyone,
@@ -111,9 +115,17 @@ def __init__(
         cpu: bool = False,
         deepspeed_plugin: DeepSpeedPlugin = None,
         rng_types: Optional[List[Union[str, RNGType]]] = None,
+        log_with: Optional[List[Union[LoggerType, str]]] = None,
         dispatch_batches: Optional[bool] = None,
         kwargs_handlers: Optional[List[KwargsHandler]] = None,
     ):
+        if log_with is not None:
+            if not isinstance(log_with, list):
+                log_with = [log_with]
+            for log_type in log_with:
+                if str(log_type) not in LoggerType:
+                    raise ValueError(f"Unsupported logging capability: {log_type}. Choose between {LoggerType.list()}")
+
         if mixed_precision is not None:
             mixed_precision = str(mixed_precision)
             if mixed_precision not in PrecisionType:
@@ -556,6 +568,30 @@ def wait_for_everyone(self):
         """
         wait_for_everyone()
 
+    def init_trackers(self, project_name: str, config: dict = None):
+        """
+        Initializes a run for all trackers stored in `self.log_with`, potentially with starting configurations
+
+        Args:
+            project_name (`str`):
+        """
+        project_location = Path(project_name)
+        self.trackers = []
+        for tracker in self.log_with:
+            if tracker.lower() == "tensorboard" and is_tensorboard_available():
+                self.trackers.append(TensorBoardTracker(project_location))
+        if config is not None:
+            for tracker in self.trackers:
+                tracker.store_init_configuration(config)
+
+    def log(self, values: dict):
+        """
+        Logs `values` to all stored trackers in `self.trackers`. Values should be a dictionary-like object containing
+        only types `int`, `float`, or `str`.
+        """
+        for tracker in self.trackers:
+            tracker.log(values)
+
     def save(self, obj, f):
         """
         Save the object passed to disk once per machine. Use in place of `torch.save`.
diff --git a/src/accelerate/tracking.py b/src/accelerate/tracking.py
new file mode 100644
index 00000000000..f8f7abccc1e
--- /dev/null
+++ b/src/accelerate/tracking.py
@@ -0,0 +1,72 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Expectation:
+# Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
+
+import logging
+from abc import ABCMeta, abstractmethod
+from pathlib import Path
+
+from torch.utils import tensorboard
+
+
+logger = logging.getLogger(__name__)
+
+
+class GeneralTracker(object, metaclass=ABCMeta):
+    """
+    A base Tracker class to be used for all logging integration implementations.
+    """
+
+    log_directory: str
+
+    @abstractmethod
+    def store_init_configuration(self, values):
+        pass
+
+    @abstractmethod
+    def log(self, values):
+        pass
+
+
+class TensorBoardTracker(GeneralTracker):
+    log_directory = "tensorboard"
+
+    def __init__(self, run_name=""):
+        self.run_name = Path(run_name)
+        self.writer = tensorboard.SummaryWriter(self.logging_dir)
+        logger.info(f"Initialized TensorBoard project {self.run_name} writing to {self.logging_dir}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    def store_init_configuration(self, values):
+        self.writer.add_hparams(values, metric_dict={})
+        self.writer.flush()
+        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+
+    # add_scalar has the option for `global_step`, should we include this?
+    # Need to see how rest of API fleshes out first with other integrations
+    def log(self, values: dict):
+        """
+        Logs `values`. Values should be a dictionary-like object containing only types `int`, `float`, or `str`.
+        """
+        for k, v in values.items():
+            if isinstance(v, (int, float)):
+                self.writer.add_scalar(k, v)
+            elif isinstance(v, str):
+                self.writer.add_text(k, v)
+        self.writer.flush()
+        logger.info("Successfully logged to TensorBoard")
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index af9299a4f96..e1329b3b2c1 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -34,6 +34,10 @@
     import torch_xla.core.xla_model as xm
 
 
+def is_tensorboard_available():
+    return importlib.util.find_spec("tensorboard") is not None or importlib.util.find_spec("tensorboardX") is not None
+
+
 def is_boto3_available():
     return importlib.util.find_spec("boto3") is not None
 
@@ -74,6 +78,11 @@ def list(cls):
         return list(map(lambda item: str(item), cls))
 
 
+class LoggerType(BaseEnum):
+    ALL = "all"
+    TENSORBOARD = "tensorboard"
+
+
 class PrecisionType(BaseEnum):
     NO = "no"
     FP16 = "fp16"


Code review comments:
diff_hunk: @@ -0,0 +1,72 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Expectation:
+# Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
+
+import logging
+from abc import ABCMeta, abstractmethod
+from pathlib import Path
+
+from torch.utils import tensorboard
path: src/accelerate/tracking.py
text: @user1:
Should be inside an `is_tensorboard_available` flag.

diff_hunk: @@ -0,0 +1,72 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Expectation:
+# Provide a project dir name, then each type of logger gets stored in project/{`logging_dir`}
+
+import logging
+from abc import ABCMeta, abstractmethod
+from pathlib import Path
+
+from torch.utils import tensorboard
+
+
+logger = logging.getLogger(__name__)
+
+
+class GeneralTracker(object, metaclass=ABCMeta):
+    """
+    A base Tracker class to be used for all logging integration implementations.
+    """
+
+    log_directory: str
+
+    @abstractmethod
+    def store_init_configuration(self, values):
+        pass
+
+    @abstractmethod
+    def log(self, values):
+        pass
+
+
+class TensorBoardTracker(GeneralTracker):
+    log_directory = "tensorboard"
+
+    def __init__(self, run_name=""):
+        self.run_name = Path(run_name)
+        self.writer = tensorboard.SummaryWriter(self.logging_dir)
+        logger.info(f"Initialized TensorBoard project {self.run_name} writing to {self.logging_dir}")
+        logger.info(
+            "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
+        )
+
+    def store_init_configuration(self, values):
+        self.writer.add_hparams(values, metric_dict={})
+        self.writer.flush()
+        logger.info("Stored initial configuration hyperparameters to TensorBoard")
+
+    # add_scalar has the option for `global_step`, should we include this?
path: src/accelerate/tracking.py
text: @user1:
Accelerate has no idea about the global step though. We could add it in an optional argument to `log`? Or add something that keeps track of the number of time the `optimizer.step` is called internally.

@author:
That's more or less what my thinking is, having it be an optional argument. I don't think there's a need for us to keep track of `optimizer.step` calls internally for this yet, but we'll see what changes as this develops

@author:
Decided to make it an optional argument. If folks seem to find it annoying then we can go ahead and track `optimizer.step` next

diff_hunk: @@ -17,10 +17,12 @@
 import sys
 import warnings
 from contextlib import contextmanager
+from pathlib import Path
 from typing import List, Optional, Union
 
 import torch
 
+from accelerate.accelerate.src.accelerate.tracking import TensorBoardTracker
path: src/accelerate/accelerator.py
text: @user1:
```suggestion
from .tracking import TensorBoardTracker
```
:scream: 

@author:
*That is what we call VSCode not being friendly and me missing it!*