added callback on hook, decoder, image logger, tried tuning

data_preprocessing.py

@@ -213,7 +213,7 @@ class LatexImageDataModule(pl.LightningDataModule):
                           pin_memory=PIN_MEMORY, num_workers=NUM_WORKERS, persistent_workers=PERSISTENT_WORKERS)
 
     def val_dataloader(self):
-        return DataLoader(self.val_dataset, batch_size=self.batch_size, collate_fn=self.collate_fn,
+        return DataLoader(self.val_dataset, batch_size=self.batch_size, shuffle=True, collate_fn=self.collate_fn,
                           pin_memory=PIN_MEMORY, num_workers=NUM_WORKERS, persistent_workers=PERSISTENT_WORKERS)
 
     def test_dataloader(self):

model.py

@@ -111,7 +111,9 @@ class Transformer(pl.LightningModule):
                  pad_idx: int,
                  dim_feedforward: int = 512,
                  dropout: float = .1,
-                 learning_rate=1e-4):
+                 learning_rate=1e-3,
+                 tex_tokenizer=None
+                 ):
         super().__init__()
 
         self.transformer = nn.Transformer(d_model=emb_size,
@@ -130,8 +132,11 @@ class Transformer(pl.LightningModule):
         self.tgt_tok_emb = TexEmbedding(emb_size, tgt_vocab_size, dropout=dropout)
         self.loss_fn = torch.nn.CrossEntropyLoss(ignore_index=pad_idx)
         self.learning_rate = learning_rate
+        self.save_hyperparameters()
+        self.tex_tokenizer = tex_tokenizer
 
-    def forward(self, src, tgt, src_mask, tgt_mask, memory_mask, src_padding_mask, tgt_padding_mask):
+    def forward(self, src, tgt, src_mask=None, tgt_mask=None, memory_mask=None, src_padding_mask=None,
+                tgt_padding_mask=None):
         src = self.src_tok_emb(src)
         tgt = self.tgt_tok_emb(tgt)
 
@@ -176,5 +181,54 @@ class Transformer(pl.LightningModule):
         return loss
 
     def configure_optimizers(self):
-        # TODO write scheduler
-        return torch.optim.Adam(self.parameters(), lr=self.learning_rate)
+        optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
+        scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=50, T_mult=1)
+        return [optimizer], [scheduler]
+
+
+class _TransformerTuner(Transformer):
+    """
+    When using trainer.tune, batches from dataloader get passed directly to forward,
+    so this subclass takes care of that
+    """
+
+    def forward(self, batch, batch_idx):
+        src = batch['images']
+        tgt = batch['tex_ids']
+        tgt_input = tgt[:, :-1]
+        tgt_output = tgt[:, 1:]
+        src_mask = None
+        tgt_mask = self.transformer.generate_square_subsequent_mask(tgt_input.shape[1]).to(self.device,
+                                                                                           torch.ByteTensor.dtype)
+        memory_mask = None
+        src_padding_mask = None
+        tgt_padding_mask = batch['tex_attention_masks'][:, :-1]
+        tgt_padding_mask = tgt_padding_mask.masked_fill(
+            tgt_padding_mask == 0, float('-inf')
+        ).masked_fill(
+            tgt_padding_mask == 1, 0
+        )
+
+        src = self.src_tok_emb(src)
+        tgt_input = self.tgt_tok_emb(tgt_input)
+        outs = self.transformer(src, tgt_input, src_mask, tgt_mask, memory_mask, src_padding_mask, tgt_padding_mask)
+        outs = self.generator(outs)
+
+        loss = self.loss_fn(einops.rearrange(outs, 'b n prob -> b prob n'), tgt_output.long())
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        return self(batch, batch_idx)
+
+
+@torch.inference_mode()
+def decode(transformer, tex_tokenizer, image):
+    tex_ids = [tex_tokenizer.token_to_id("[CLS]")]
+    while tex_ids[-1] != tex_tokenizer.token_to_id("[SEP]") and len(tex_ids) < 30:
+        src = einops.rearrange(image, "c h w -> () c h w")
+        tgt = torch.tensor([tex_ids], device=transformer.device, dtype=torch.float32)
+        outs = transformer(src, tgt)
+        next_id = outs[:, -1].argmax(dim=1).item()
+        tex_ids.append(next_id)
+    tex = tex_tokenizer.decode(tex_ids, skip_special_tokens=True)
+    return tex

train.py

@@ -1,26 +1,32 @@
 from data_generator import generate_data
 from data_preprocessing import LatexImageDataModule, IMAGE_WIDTH, IMAGE_HEIGHT
-from model import Transformer
+from model import Transformer, _TransformerTuner
+from utils import LogImageTexCallback
 
 import argparse
-from pytorch_lightning.loggers import WandbLogger
+from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger
 from pytorch_lightning import Trainer, seed_everything
 import torch
+import wandb
 
-DATASET_PATH = 'resources/dataset.pt'
+DATASET_PATH = "resources/dataset.pt"
+TRAINER_DIR = "resources/pl_trainer_checkpoints"
+TUNER_DIR = "resources/pl_tuner_checkpoints"
+TRAINER_STRATEGY = "ddp"
+BEST_MODEL_CHECKPOINT = "best_model.ckpt"
 
 
 def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument(
-        "-m", "-max-epochs", help="limit the number of training epochs", type=int, dest='max_epochs'
+        "-m", "-max-epochs", help="limit the number of training epochs", type=int, dest="max_epochs"
     )
     parser.add_argument(
         "-n", "-new-dataset", help="clear old dataset and generate provided number of new examples", type=int,
         dest="new_dataset"
     )
     parser.add_argument(
-        "-g", "-gpus", help=f"number of gpus to train on in range 0..{torch.cuda.device_count()}",
+        "-g", "-gpus", metavar="GPUS", help="ids of gpus to train on, if not provided then trains on cpu", nargs="+",
         type=int, dest="gpus", choices=list(range(torch.cuda.device_count())),
     )
     parser.add_argument(
@@ -31,6 +37,10 @@ def parse_args():
         "-d", "-deterministic", help="whether to seed all rngs for reproducibility, default False", default=False,
         action="store_true", dest="deterministic"
     )
+    # parser.add_argument(
+    #     "-t", "-tune", help="whether to tune model for batch size before training, default False", default=False,
+    #     action="store_true", dest="tune"
+    # )
 
     args = parser.parse_args()
     return args
@@ -52,17 +62,21 @@ def main():
     # TODO: log images, accuracy?, update python, write own transformer, add checkpoints, lr scheduler,
     #  determine when trainer doesnt hang(when  single gpu,ddp, num_workers=0)
 
-    logger = WandbLogger(f"img2tex", version='0') if args.log else False
+    if args.log:
+        logger = WandbLogger(f"img2tex", log_model=True)
+        callbacks = [LogImageTexCallback(logger, datamodule.tex_tokenizer)]
+    else:
+        logger = None
+        callbacks = []
 
     trainer = Trainer(max_epochs=args.max_epochs,
-                      accelerator='gpu' if args.gpus else 'cpu',
+                      accelerator="cpu" if args.gpus is None else "gpu",
                       gpus=args.gpus,
                       logger=logger,
-                      strategy='ddp',
-                      auto_scale_batch_size="power",
-                      auto_lr_find=True,
-                      auto_select_gpus=True,
-                      enable_progress_bar=True
+                      strategy=TRAINER_STRATEGY,
+                      enable_progress_bar=True,
+                      default_root_dir=TRAINER_DIR,
+                      callbacks=callbacks,
                       )
 
     transformer = Transformer(num_encoder_layers=3,
@@ -77,11 +91,31 @@ def main():
                               dropout=0.1
                               )
 
-    trainer.tune(transformer, datamodule=datamodule)
+    # dl = datamodule.train_dataloader()
+    # b = next(iter(dl))
+    # image=b['images'][0]
+    # tex = decode(transformer, datamodule.tex_tokenizer, image)
+    # print(tex)
+
+    # if args.new_dataset:
+    #     datamodule.batch_size = 1
+    #     transformer_for_tuning = TransformerTuner(**transformer.hparams).cuda()
+    #     tuner = Trainer(accelerator="gpu" if args.gpus else "cpu",
+    #                     gpus=args.gpus,
+    #                     strategy=TRAINER_STRATEGY,
+    #                     enable_progress_bar=True,
+    #                     enable_checkpointing=False,
+    #                     auto_scale_batch_size=True,
+    #                     num_sanity_val_steps=0,
+    #                     logger=False
+    #                     )
+    #     tuner.tune(transformer_for_tuning, datamodule=datamodule)
+    #     torch.save(datamodule, DATASET_PATH)
+
     trainer.fit(transformer, datamodule=datamodule)
     trainer.test(datamodule=datamodule)
-    trainer.save_checkpoint("best_model.ckpt")
+    trainer.save_checkpoint(BEST_MODEL_CHECKPOINT)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()

utils.py

@@ -0,0 +1,21 @@
+import torch
+from pytorch_lightning.callbacks import Callback
+from model import decode
+
+from torchvision import transforms
+
+
+class LogImageTexCallback(Callback):
+    def __init__(self, logger, tex_tokenizer):
+        self.logger = logger
+        self.tex_tokenizer = tex_tokenizer
+        self.tensor_to_PIL = transforms.ToPILImage()
+
+    def on_validation_batch_start(self, trainer, transformer, batch, batch_idx, dataloader_idx):
+        if batch_idx != 0 or dataloader_idx != 0:
+            return
+        image = batch['images'][0]
+        tex_predicted = decode(transformer, self.tex_tokenizer, image)
+        image = self.tensor_to_PIL(image)
+        tex_true = self.tex_tokenizer.decode(list(batch['tex_ids'][0].to('cpu', torch.int)), skip_special_tokens=True)
+        self.logger.log_image(key="samples", images=[image], caption=[f"True {tex_true}\n Predicted{tex_predicted}"])