add wandb logger

Author: timothyxp

.gitignore

@@ -104,6 +104,7 @@ ENV/
 data/
 input/
 saved/
+wandb/
 
 # editor, os cache directory
 .vscode/

hw_asr/base/base_trainer.py

@@ -4,7 +4,7 @@
 from numpy import inf
 
 from hw_asr.base import BaseModel
-from hw_asr.logger import TensorboardWriter
+from hw_asr.logger import get_visualizer
 
 
 class BaseTrainer:
@@ -48,8 +48,8 @@ def __init__(self, model: BaseModel, criterion, metrics, optimizer, config, devi
         self.checkpoint_dir = config.save_dir
 
         # setup visualization writer instance
-        self.writer = TensorboardWriter(
-            config.log_dir, self.logger, cfg_trainer["tensorboard"]
+        self.writer = get_visualizer(
+            config, self.logger, cfg_trainer["visualize"]
         )
 
         if config.resume is not None:

hw_asr/config.json

@@ -91,7 +91,8 @@
     "verbosity": 2,
     "monitor": "min val_loss",
     "early_stop": 100,
-    "tensorboard": true,
+    "visualize": "wandb",
+    "wandb_project": "asr_project",
     "len_epoch": 100,
     "grad_norm_clip": 10
   }

hw_asr/logger/tensorboard.py

@@ -0,0 +1,87 @@
+import importlib
+from datetime import datetime
+
+
+class TensorboardWriter:
+    def __init__(self, log_dir, logger, enabled):
+        self.writer = None
+        self.selected_module = ""
+
+        if enabled:
+            log_dir = str(log_dir)
+
+            # Retrieve vizualization writer.
+            succeeded = False
+            for module in ["torch.utils.tensorboard", "tensorboardX"]:
+                try:
+                    self.writer = importlib.import_module(module).SummaryWriter(log_dir)
+                    succeeded = True
+                    break
+                except ImportError:
+                    succeeded = False
+                self.selected_module = module
+
+            if not succeeded:
+                message = (
+                    "Warning: visualization (Tensorboard) is configured to use, but currently not installed on "
+                    "this machine. Please install TensorboardX with 'pip install tensorboardx', upgrade PyTorch to "
+                    "version >= 1.1 to use 'torch.utils.tensorboard' or turn off the option in the 'config.json' file."
+                )
+                logger.warning(message)
+
+        self.step = 0
+        self.mode = ""
+
+        self.tb_writer_ftns = {
+            "add_scalar",
+            "add_scalars",
+            "add_image",
+            "add_images",
+            "add_audio",
+            "add_text",
+            "add_histogram",
+            "add_pr_curve",
+            "add_embedding",
+        }
+        self.tag_mode_exceptions = {"add_histogram", "add_embedding"}
+        self.timer = datetime.now()
+
+    def set_step(self, step, mode="train"):
+        self.mode = mode
+        self.step = step
+        if step == 0:
+            self.timer = datetime.now()
+        else:
+            duration = datetime.now() - self.timer
+            self.add_scalar("steps_per_sec", 1 / duration.total_seconds())
+            self.timer = datetime.now()
+
+    def __getattr__(self, name):
+        """
+        If visualization is configured to use:
+            return add_data() methods of tensorboard with additional information (step, tag) added.
+        Otherwise:
+            return a blank function handle that does nothing
+        """
+        if name in self.tb_writer_ftns:
+            add_data = getattr(self.writer, name, None)
+
+            def wrapper(tag, data, *args, **kwargs):
+                if add_data is not None:
+                    # add mode(train/valid) tag
+                    if name not in self.tag_mode_exceptions:
+                        tag = "{}/{}".format(tag, self.mode)
+                    add_data(tag, data, self.step, *args, **kwargs)
+
+            return wrapper
+        else:
+            # default action for returning methods defined in this class, set_step() for instance.
+            try:
+                attr = object.__getattr__(name)
+            except AttributeError:
+                raise AttributeError(
+                    "type object '{}' has no attribute '{}'".format(
+                        self.selected_module, name
+                    )
+                )
+            return attr

hw_asr/logger/visualization.py

@@ -1,87 +1,13 @@
-import importlib
-from datetime import datetime
+from .tensorboard import TensorboardWriter
+from .wandb import WanDBdWriter
 
 
-class TensorboardWriter:
-    def __init__(self, log_dir, logger, enabled):
-        self.writer = None
-        self.selected_module = ""
+def get_visualizer(config, logger, type):
+    if type == "tensorboard":
+        return TensorboardWriter(config.log_dir, logger, True)
 
-        if enabled:
-            log_dir = str(log_dir)
+    if type == 'wandb':
+        return WanDBdWriter(config, logger)
 
-            # Retrieve vizualization writer.
-            succeeded = False
-            for module in ["torch.utils.tensorboard", "tensorboardX"]:
-                try:
-                    self.writer = importlib.import_module(module).SummaryWriter(log_dir)
-                    succeeded = True
-                    break
-                except ImportError:
-                    succeeded = False
-                self.selected_module = module
+    return None
 
-            if not succeeded:
-                message = (
-                    "Warning: visualization (Tensorboard) is configured to use, but currently not installed on "
-                    "this machine. Please install TensorboardX with 'pip install tensorboardx', upgrade PyTorch to "
-                    "version >= 1.1 to use 'torch.utils.tensorboard' or turn off the option in the 'config.json' file."
-                )
-                logger.warning(message)
-
-        self.step = 0
-        self.mode = ""
-
-        self.tb_writer_ftns = {
-            "add_scalar",
-            "add_scalars",
-            "add_image",
-            "add_images",
-            "add_audio",
-            "add_text",
-            "add_histogram",
-            "add_pr_curve",
-            "add_embedding",
-        }
-        self.tag_mode_exceptions = {"add_histogram", "add_embedding"}
-        self.timer = datetime.now()
-
-    def set_step(self, step, mode="train"):
-        self.mode = mode
-        self.step = step
-        if step == 0:
-            self.timer = datetime.now()
-        else:
-            duration = datetime.now() - self.timer
-            self.add_scalar("steps_per_sec", 1 / duration.total_seconds())
-            self.timer = datetime.now()
-
-    def __getattr__(self, name):
-        """
-        If visualization is configured to use:
-            return add_data() methods of tensorboard with additional information (step, tag) added.
-        Otherwise:
-            return a blank function handle that does nothing
-        """
-        if name in self.tb_writer_ftns:
-            add_data = getattr(self.writer, name, None)
-
-            def wrapper(tag, data, *args, **kwargs):
-                if add_data is not None:
-                    # add mode(train/valid) tag
-                    if name not in self.tag_mode_exceptions:
-                        tag = "{}/{}".format(tag, self.mode)
-                    add_data(tag, data, self.step, *args, **kwargs)
-
-            return wrapper
-        else:
-            # default action for returning methods defined in this class, set_step() for instance.
-            try:
-                attr = object.__getattr__(name)
-            except AttributeError:
-                raise AttributeError(
-                    "type object '{}' has no attribute '{}'".format(
-                        self.selected_module, name
-                    )
-                )
-            return attr

hw_asr/logger/wandb.py

@@ -0,0 +1,79 @@
+from datetime import datetime
+
+
+class WanDBdWriter:
+    def __init__(self, config, logger):
+        self.writer = None
+        self.selected_module = ""
+
+        try:
+            import wandb
+            wandb.login()
+
+            if config['trainer'].get('wandb_project') is None:
+                raise ValueError("please specify project name for wandb")
+
+            wandb.init(
+                project=config['trainer'].get('wandb_project'),
+                config=config.config
+            )
+            self.wandb = wandb
+
+        except ImportError:
+            logger.warning("For use wandb install it via \n\t pip install wandb")
+
+        self.step = 0
+        self.mode = ""
+        self.timer = datetime.now()
+
+    def set_step(self, step, mode="train"):
+        self.mode = mode
+        self.step = step
+        if step == 0:
+            self.timer = datetime.now()
+        else:
+            duration = datetime.now() - self.timer
+            self.add_scalar("steps_per_sec", 1 / duration.total_seconds())
+            self.timer = datetime.now()
+
+    def scalar_name(self, scalar_name):
+        return f"{scalar_name}_{self.mode}"
+
+    def add_scalar(self, scalar_name, scalar):
+        self.wandb.log({
+            self.scalar_name(scalar_name): scalar,
+        }, step=self.step)
+
+    def add_scalars(self, tag, scalars):
+        self.wandb.log({
+            **{f"{scalar_name}_{tag}_{self.mode}": scalar for scalar_name, scalar in scalars.items()}
+        }, step=self.step)
+
+    def add_image(self, scalar_name, image):
+        self.wandb.log({
+            self.scalar_name(scalar_name): self.wandb.Image(image)
+        }, step=self.step)
+
+    def add_audio(self, scalar_name, audio):
+        self.wandb.log({
+            self.scalar_name(scalar_name): self.wandb.Audio(audio)
+        }, step=self.step)
+
+    def add_text(self, scalar_name, text):
+        self.wandb.log({
+            self.scalar_name(scalar_name): self.wandb.Html(text)
+        }, step=self.step)
+
+    def add_histogram(self, scalar_name, hist, bins=None):
+        hist = hist.detach().cpu().numpy()
+        hist = self.wandb.Histogram(hist, num_bins=bins)
+
+        self.wandb.log({
+            self.scalar_name(scalar_name): hist
+        }, step=self.step)
+
+    def add_pr_curve(self, scalar_name, scalar):
+        raise NotImplementedError()
+
+    def add_embedding(self, scalar_name, scalar):
+        raise NotImplementedError()

hw_asr/trainer/trainer.py

@@ -86,8 +86,10 @@ def _train_iteration(self, batch: dict, epoch: int, batch_num: int):
         batch["log_probs_length"] = self.model.transform_input_lengths(
             batch["spectrogram_length"]
         )
+
         loss = self.criterion(**batch)
         loss.backward()
+
         self._clip_grad_norm()
         self.optimizer.step()
 

requirements.txt

@@ -3,6 +3,8 @@ torchvision
 numpy
 tqdm
 tensorboard
+matplotlib
+
 
 pandas
 speechbrain~=0.5.9