wandb support

Author: lukas-blecher

models.py

@@ -6,7 +6,6 @@
 from x_transformers.autoregressive_wrapper import AutoregressiveWrapper
 from einops import rearrange, repeat
 
-
 class ViTransformerWrapper(nn.Module):
     def __init__(
         self,
@@ -90,6 +89,9 @@ def get_model(args):
                 heads=args.heads,
                 cross_attend=True
             )),
-        pad_value=args.pad_token_id
+        pad_value=args.pad_token
     ).to(args.device)
+    if args.wandb:
+        import wandb
+        wandb.watch((encoder, decoder))
     return Model(encoder, decoder, args)

settings/default.yaml

@@ -2,13 +2,9 @@
 data: "dataset/data/dataset.pkl"
 output_path: "outputs"
 model_path: "checkpoints"
+save_freq: 5 # save every nth epoch
 name: "pix2tex"
 
-# Token ids
-pad_token: 0
-bos_token: 1
-eos_token: 2
-
 # Training parameters
 epochs: 10
 batchsize: 8
@@ -29,3 +25,14 @@ num_layers: 4
 heads: 8
 num_tokens: 8000
 max_seq_len: 512
+
+# Other
+seed: 42
+id: null
+sample_freq: 50
+debug: True
+
+# Token ids
+pad_token: 0
+bos_token: 1
+eos_token: 2

train.py

@@ -14,13 +14,15 @@
 
 from dataset.dataset import Im2LatexDataset
 from models import get_model
+from utils import *
 
 
 def train(args):
     dataloader = Im2LatexDataset().load(args.data)
     dataloader.update(args)
     device = args.device
-    args.pad_token_id = dataloader.pad_token_id
+    os.makedirs(args.model_path, exist_ok=True)
+
     model = get_model(args)
     encoder, decoder = model.encoder, model.decoder
     opt = optim.Adam(model.parameters(), args.lr)
@@ -36,10 +38,18 @@ def train(args):
             torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
             opt.step()
             dset.set_description('Loss: %.4f' % loss.item())
-            if i % 15 == 0:
-                print(''.join(dataloader.tokenizer.decode(decoder.generate(torch.LongTensor([dataloader.bos_token_id]).to(
-                    device), args.max_seq_len, eos_token=dataloader.eos_token_id, context=encoded[:1])[:-1]).split(' ')).replace('Ġ', ' ').strip())
-                print(dataloader.pairs[dataloader.i][0][0])
+            if args.wandb:
+                wandb.log({'train/loss': loss.item()})
+            if i % args.sample_freq == 0:
+                pred = ''.join(dataloader.tokenizer.decode(decoder.generate(torch.LongTensor([dataloader.bos_token_id]).to(
+                    device), args.max_seq_len, eos_token=dataloader.eos_token_id, context=encoded[:1])[:-1]).split(' ')).replace('Ġ', ' ').strip()
+                truth = dataloader.pairs[dataloader.i][0][0]
+                if args.wandb:
+                    table = wandb.Table(columns=["Truth", "Prediction"])
+                    table.add_data(tuth, pred)
+                    wandb.log({"test/examples": table})
+        if (e+1) % args.save_freq == 0:
+            torch.save(model.parameters(), os.path.join(args.model_path, '%s_e%02d' % (args.name, e+1)))
 
 
 if __name__ == '__main__':
@@ -54,6 +64,12 @@ def train(args):
     with parsed_args.config as f:
         params = yaml.load(f, Loader=yaml.FullLoader)
     args = Munch(params)
+    args.wandb = not parsed_args.debug and not args.debug
     logging.getLogger().setLevel(logging.DEBUG if parsed_args.debug else logging.WARNING)
     args.device = torch.device('cuda' if torch.cuda.is_available() and not parsed_args.no_cuda else 'cpu')
+    seed_everything(args.seed)
+    if args.wandb:
+        if not parsed_args.resume:
+            args.id = wandb.util.generate_id()
+        wandb.init(config=dict(args), resume='allow', name=args.name, id=args.id)
     train(args)

utils/__init__.py

@@ -0,0 +1 @@
+from utils.utils import *

utils/utils.py

@@ -0,0 +1,13 @@
+def seed_everything(seed: int):
+    import random
+    import os
+    import numpy as np
+    import torch
+
+    random.seed(seed)
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = True