add TensorBoardLogger to transformers integration (#912)

Author: bfineran

src/sparseml/pytorch/utils/logger.py

@@ -778,6 +778,16 @@ def __len__(self):
     def __iter__(self):
         return iter(self.loggers)
 
+    def add_logger(self, logger: BaseLogger):
+        """
+        add a BaseLogger implementation to the loggers of this manager
+
+        :param logger: logger object to add
+        """
+        if not isinstance(logger, BaseLogger):
+            raise ValueError(f"logger {type(logger)} must be of type BaseLogger")
+        self._loggers.append(logger)
+
     def log_ready(self, epoch, last_log_epoch):
         """
         Check if there is a logger that is ready to accept a log

src/sparseml/transformers/sparsification/trainer.py

@@ -32,6 +32,7 @@
 from transformers import Trainer as TransformersTrainer
 from transformers import TrainerCallback, TrainerControl, TrainingArguments
 from transformers.file_utils import WEIGHTS_NAME
+from transformers.integrations import TensorBoardCallback
 from transformers.trainer_callback import TrainerState
 from transformers.trainer_utils import get_last_checkpoint
 
@@ -40,6 +41,7 @@
     GradSampler,
     LoggerManager,
     ModuleSparsificationInfo,
+    TensorBoardLogger,
     WANDBLogger,
 )
 from sparseml.transformers.utils import SparseAutoModel
@@ -154,7 +156,7 @@ def __init__(
         self.criterion = torch.nn.CrossEntropyLoss()
         self.callback_disable_fp16 = DisableHalfPrecisionCallback(self)
         self.callback_handler.add_callback(self.callback_disable_fp16)
-
+        self._add_tensorboard_logger_if_available()
         self.grad_sampler = GradSampler(
             self._mfac_data_loader(), self._mfac_loss_function
         )
@@ -263,7 +265,6 @@ def create_optimizer(self):
         self.manager_steps_per_epoch = math.ceil(
             len(self.train_dataset) / total_batch_size
         )
-
         if hasattr(self, "scaler"):
             wrap_optim_key = "scaler"
             self.scaler = self.manager.modify(
@@ -702,6 +703,24 @@ def _mfac_loss_function(self, model_outputs, loss_target):
             )
         return loss
 
+    def _add_tensorboard_logger_if_available(self):
+        tensorboard_callback = None
+        for callback in self.callback_handler.callbacks:
+            if isinstance(callback, TensorBoardCallback):
+                tensorboard_callback = callback
+                break
+        if tensorboard_callback is None:
+            return
+
+        if tensorboard_callback.tb_writer is None:
+            tensorboard_callback._init_summary_writer(
+                self.args, log_dir=self.args.logging_dir
+            )
+
+        self.logger_manager.add_logger(
+            TensorBoardLogger(writer=tensorboard_callback.tb_writer)
+        )
+
 
 class TrainerInterface(RecipeManagerTrainerInterface):
     """