noam lr scheduler, shared weight between embedding and generator

data_preprocessing.py

@@ -115,8 +115,10 @@ class RandomizeImageTransform(object):
     """Standardize image and randomly augment"""
 
     def __init__(self, width=IMAGE_WIDTH, height=IMAGE_HEIGHT, random_magnitude=5):
+        assert random_magnitude > 0
+        eps = 0.01
         self.transform = T.Compose((
-            T.ColorJitter(brightness=random_magnitude / 10),
+            T.ColorJitter(brightness=((1 - eps) / (random_magnitude + eps), 1 - eps)),
             T.Resize(height),
             T.Grayscale(),
             T.functional.invert,
@@ -184,12 +186,12 @@ class LatexImageDataModule(pl.LightningDataModule):
         )
         self.val_dataset = TexImageDataset(
             root_dir=DATA_DIR,
-            image_transform=StandardizeImageTransform(),
+            image_transform=RandomizeImageTransform(),
             tex_transform=ExtractEquationFromTexTransform()
         )
         self.test_dataset = TexImageDataset(
             root_dir=DATA_DIR,
-            image_transform=StandardizeImageTransform(),
+            image_transform=RandomizeImageTransform(),
             tex_transform=ExtractEquationFromTexTransform()
         )
         train_indices, val_indices, test_indices = self.train_val_test_split(len(self.train_dataset))
@@ -215,7 +217,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, shuffle=True, collate_fn=self.collate_fn,
+        return DataLoader(self.val_dataset, batch_size=self.batch_size, collate_fn=self.collate_fn,
                           pin_memory=PIN_MEMORY, num_workers=NUM_WORKERS, persistent_workers=PERSISTENT_WORKERS)
 
     def test_dataloader(self):

model.py

@@ -41,7 +41,7 @@ class AddPositionalEncoding(nn.Module):
 class ImageEmbedding(nn.Module):
     """Reshape image into patches and project into given dimension"""
 
-    def __init__(self, d_model, input_width, input_height, patch_size=16, dropout=.1):
+    def __init__(self, d_model, input_width, input_height, patch_size, dropout):
         super().__init__()
         assert input_width % patch_size == 0 and input_height % patch_size == 0, \
             "Cannot split image in patches"
@@ -64,7 +64,7 @@ class ImageEmbedding(nn.Module):
 
 
 class TexEmbedding(nn.Module):
-    def __init__(self, d_model: int, vocab_size: int, dropout: float = .1):
+    def __init__(self, d_model: int, vocab_size: int, dropout: float):
         super().__init__()
         self.embedding = nn.Embedding(vocab_size, d_model)
         self.add_positional_encoding = AddPositionalEncoding(d_model)
@@ -85,7 +85,7 @@ class ImageEncoder(nn.Module):
 
     def __init__(self, image_width, image_height, d_model, num_layers=8):
         super().__init__()
-        image_embedding = ImageEmbedding(d_model, image_width, image_height)
+        image_embedding = ImageEmbedding(d_model, image_width, image_height, patch_size=16, dropout=.1)
         encoder_layer = nn.TransformerEncoderLayer(
             d_model=d_model,
             nhead=8,
@@ -111,7 +111,6 @@ class Transformer(pl.LightningModule):
                  pad_idx: int,
                  dim_feedforward: int = 512,
                  dropout: float = .1,
-                 learning_rate: float = 1e-3
                  ):
         super().__init__()
 
@@ -126,11 +125,13 @@ class Transformer(pl.LightningModule):
             if p.dim() > 1:
                 nn.init.xavier_uniform_(p)
 
-        self.generator = nn.Linear(emb_size, tgt_vocab_size)
-        self.src_tok_emb = ImageEmbedding(emb_size, image_width, image_height, dropout=dropout)
+        self.d_model = emb_size
+        self.src_tok_emb = ImageEmbedding(emb_size, image_width, image_height, patch_size=16, dropout=dropout)
         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.generator = nn.Linear(emb_size, tgt_vocab_size)
+        # Make embedding and generator share weight because they do the same thing
+        self.tgt_tok_emb.embedding.weight = self.generator.weight
+        self.loss_fn = torch.nn.CrossEntropyLoss(ignore_index=pad_idx, label_smoothing=.1)
         self.save_hyperparameters()
 
     def forward(self, src, tgt, src_mask=None, tgt_mask=None, memory_mask=None, src_padding_mask=None,
@@ -147,16 +148,10 @@ class Transformer(pl.LightningModule):
         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)
+        tgt_mask = self.transformer.generate_square_subsequent_mask(tgt_input.shape[1]).to(self.device, torch.bool)
         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
-        )
+        tgt_padding_mask = torch.logical_not(batch['tex_attention_masks'][:, :-1])
 
         outs = self(src, tgt_input, src_mask, tgt_mask, memory_mask, src_padding_mask, tgt_padding_mask)
         loss = self.loss_fn(einops.rearrange(outs, 'b n prob -> b prob n'), tgt_output.long())
@@ -179,58 +174,22 @@ class Transformer(pl.LightningModule):
         return loss
 
     def configure_optimizers(self):
-        optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
-        scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=50, T_mult=1)
+        optimizer = torch.optim.Adam(self.parameters(), lr=1, betas=(0.9, 0.98), eps=1e-9)
+        scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, NoamLRLambda(self.d_model))
         return [optimizer], [scheduler]
 
-    # def configure_optimizers(self):
-    #     optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
-    #     return optimizer
-
-
-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
+class NoamLRLambda:
+    def __init__(self, d_model, factor=1, warmup=4000):
+        """
+        :param d_model: size of hidden model dimension
+        :param factor: multiplicative factor
+        :param warmup: number of warmup steps
+        """
+        self.d_model = d_model
+        self.factor = factor
+        self.warmup = warmup
 
-    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
+    def __call__(self, step):
+        step += 1
+        return self.factor * self.d_model ** (-0.5) * min(step ** (-0.5), step * self.warmup ** (-1.5))

train.py

@@ -1,18 +1,17 @@
 from data_generator import generate_data
 from data_preprocessing import LatexImageDataModule, IMAGE_WIDTH, IMAGE_HEIGHT
-from model import Transformer, _TransformerTuner
+from model import Transformer
 from utils import LogImageTexCallback
 
 import argparse
 from pytorch_lightning.callbacks import LearningRateMonitor
-from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger
+from pytorch_lightning.loggers import WandbLogger
 from pytorch_lightning import Trainer, seed_everything
 import torch
 
 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"
 
 
@@ -46,6 +45,10 @@ def parse_args():
     return args
 
 
+# TODO: update python, maybe model doesnt train bc of ignore special index in CrossEntropyLoss?
+# crop image, adjust brightness, lr warmup?, make tex tokens always decodable,
+# take loss that doesn't punish so much for offsets, take a look at weights,
+
 def main():
     args = parse_args()
 
@@ -58,8 +61,6 @@ def main():
         torch.save(datamodule, DATASET_PATH)
     else:
         datamodule = torch.load(DATASET_PATH)
-    # 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)
 
     if args.log:
         logger = WandbLogger(f"img2tex", log_model=True)
@@ -75,7 +76,7 @@ def main():
                       accelerator="cpu" if args.gpus is None else "gpu",
                       gpus=args.gpus,
                       logger=logger,
-                      strategy=TRAINER_STRATEGY,
+                      strategy="ddp",
                       enable_progress_bar=True,
                       default_root_dir=TRAINER_DIR,
                       callbacks=callbacks,
@@ -91,24 +92,8 @@ def main():
                               pad_idx=datamodule.tex_tokenizer.token_to_id("[PAD]"),
                               dim_feedforward=512,
                               dropout=0.1,
-                              learning_rate=1e-3
                               )
 
-    # 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_CHECKPOINT)

utils.py

@@ -1,7 +1,9 @@
-import torch
-from pytorch_lightning.callbacks import Callback
-from model import decode
+from model import Transformer
 
+import einops
+import random
+from pytorch_lightning.callbacks import Callback
+import torch
 from torchvision import transforms
 
 
@@ -14,8 +16,75 @@ class LogImageTexCallback(Callback):
     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)
+        sample_id = random.randint(0, len(batch['images']) - 1)
+        image = batch['images'][sample_id]
+        tex_predicted, tex_ids = 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}"])
+        tex_true = self.tex_tokenizer.decode(list(batch['tex_ids'][sample_id].to('cpu', torch.int)),
+                                             skip_special_tokens=True)
+        self.logger.log_image(key="samples", images=[image],
+                              caption=[f"True: {tex_true}\nPredicted: {tex_predicted}\nIds: {tex_ids}"])
+
+# 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)
+
+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]")]
+    src = einops.rearrange(image, "c h w -> () c h w")
+    while tex_ids[-1] != tex_tokenizer.token_to_id("[SEP]") and len(tex_ids) < 30:
+        tgt = torch.tensor([tex_ids], device=transformer.device, dtype=torch.float)
+        tgt_mask = transformer.transformer.generate_square_subsequent_mask(tgt.shape[1]).to(transformer.device,
+                                                                                            torch.bool)
+        outs = transformer(src, tgt, src_mask=None, tgt_mask=tgt_mask)
+        outs = einops.rearrange(outs, 'b n prob -> b prob n')
+        next_id = outs[0, :, -1].argmax().item()
+        tex_ids.append(next_id)
+    tex = tex_tokenizer.decode(tex_ids, skip_special_tokens=True)
+    return tex, tex_ids