experiment tracker agnostic

dalle2_pytorch/trackers.py

@@ -0,0 +1,49 @@
+import os
+import torch
+from torch import nn
+
+# helper functions
+
+def exists(val):
+    return val is not None
+
+# base class
+
+class BaseTracker(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def init(self, config, **kwargs):
+        raise NotImplementedError
+
+    def log(self, log, **kwargs):
+        raise NotImplementedError
+
+# basic stdout class
+
+class ConsoleTracker(BaseTracker):
+    def init(self, **config):
+        print(config)
+
+    def log(self, log, **kwargs):
+        print(log)
+
+# basic wandb class
+
+class WandbTracker(BaseTracker):
+    def __init__(self):
+        super().__init__()
+        try:
+            import wandb
+        except ImportError as e:
+            print('`pip install wandb` to use the wandb experiment tracker')
+            raise e
+
+        os.environ["WANDB_SILENT"] = "true"
+        self.wandb = wandb
+
+    def init(self, **config):
+        self.wandb.init(**config)
+
+    def log(self, log, **kwargs):
+        self.wandb.log(log, **kwargs)

dalle2_pytorch/train.py

@@ -228,6 +228,8 @@ class DiffusionPriorTrainer(nn.Module):
 
         self.max_grad_norm = max_grad_norm
 
+        self.register_buffer('step', torch.tensor([0.]))
+
     def update(self):
         if exists(self.max_grad_norm):
             self.scaler.unscale_(self.optimizer)
@@ -240,6 +242,8 @@ class DiffusionPriorTrainer(nn.Module):
         if self.use_ema:
             self.ema_diffusion_prior.update()
 
+        self.step += 1
+
     @torch.inference_mode()
     def p_sample_loop(self, *args, **kwargs):
         return self.ema_diffusion_prior.ema_model.p_sample_loop(*args, **kwargs)
@@ -328,6 +332,8 @@ class DecoderTrainer(nn.Module):
 
         self.max_grad_norm = max_grad_norm
 
+        self.register_buffer('step', torch.tensor([0.]))
+
     @property
     def unets(self):
         return nn.ModuleList([ema.ema_model for ema in self.ema_unets])
@@ -358,6 +364,8 @@ class DecoderTrainer(nn.Module):
             ema_unet = self.ema_unets[index]
             ema_unet.update()
 
+        self.step += 1
+
     @torch.no_grad()
     def sample(self, *args, **kwargs):
         if self.use_ema: