beam search

constants.py

@@ -5,6 +5,7 @@ DATA_DIR = "data"
 LATEX_PATH = "resources/latex.json"
 TRAINER_DIR = "resources/trainer"
 TOKENIZER_PATH = "resources/tokenizer.pt"
+DATAMODULE_PATH = "resources/datamodule.pt"
 
 NUM_DATALOADER_WORKERS = 4
 PERSISTENT_WORKERS = True  # whether to shut down workers at the end of epoch

data_preprocessing.py

@@ -119,7 +119,6 @@ def generate_tex_tokenizer(dataloader):
 
     texs = list(tqdm.tqdm((batch['tex'] for batch in dataloader), "Training tokenizer", total=len(dataloader)))
 
-    os.environ['TOKENIZERS_PARALLELISM'] = 'false'
     tokenizer = tokenizers.Tokenizer(tokenizers.models.BPE(unk_token="[UNK]"))
     tokenizer_trainer = tokenizers.trainers.BpeTrainer(
         special_tokens=["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]"]
@@ -150,10 +149,10 @@ class LatexImageDataModule(pl.LightningDataModule):
         self.batch_size = batch_size
         self.save_hyperparameters()
 
-    def prepare_data(self):
+    def train_tokenizer(self):
         tokenizer = generate_tex_tokenizer(DataLoader(self.train_dataset, batch_size=32, num_workers=16))
-        print(f"Vocabulary size: {tokenizer.get_vocab_size()}")
         torch.save(tokenizer, TOKENIZER_PATH)
+        return tokenizer
 
     def _shared_dataloader(self, dataset, **kwargs):
         tex_tokenizer = torch.load(TOKENIZER_PATH)

model.py

@@ -130,12 +130,24 @@ class Transformer(pl.LightningModule):
 
     def forward(self, src, tgt, src_mask=None, tgt_mask=None, memory_mask=None, src_padding_mask=None,
                 tgt_padding_mask=None):
+        """The positions of masks with ``True``
+            are not allowed to attend while ``False`` values will be unchanged.
+        The positions of padding masks with the
+            value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged."""
         src = self.src_tok_emb(src)
         tgt = self.tgt_tok_emb(tgt)
-
         outs = self.transformer(src, tgt, src_mask, tgt_mask, memory_mask, src_padding_mask, tgt_padding_mask)
         return self.generator(outs)
 
+    def encode(self, src, src_mask=None, src_padding_mask=None):
+        src = self.src_tok_emb(src)
+        return self.transformer.encoder(src, src_mask, src_padding_mask)
+
+    def decode(self, tgt, memory=None, tgt_mask=None, memory_mask=None, tgt_padding_mask=None):
+        tgt = self.tgt_tok_emb(tgt)
+        outs = self.transformer.decoder(tgt, memory, tgt_mask, memory_mask, tgt_padding_mask)
+        return self.generator(outs)
+
     def _shared_step(self, batch):
         src = batch['images']
         tgt = batch['tex_ids']

train.py

@@ -1,4 +1,4 @@
-from constants import TRAINER_DIR
+from constants import TRAINER_DIR, TOKENIZER_PATH, DATAMODULE_PATH
 from data_preprocessing import LatexImageDataModule
 from model import Transformer
 from utils import LogImageTexCallback
@@ -11,77 +11,105 @@ from pytorch_lightning import Trainer
 import torch
 
 
-# TODO: update python, maybe model doesnt train bc of ignore special index in CrossEntropyLoss?
-# crop image, adjust brightness, make tex tokens always decodable,
-# save only datamodule state?, ensemble last checkpoints, early stopping
+# TODO: update python, make tex tokens always decodable, ensemble last checkpoints,
+#  clear checkpoint data build full dataset, train export model to torchscript write spaces interface
+
+def check_setup():
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+    if not os.path.isfile(DATAMODULE_PATH):
+        datamodule = LatexImageDataModule(image_width=1024, image_height=128, batch_size=16, random_magnitude=5)
+        torch.save(datamodule, DATAMODULE_PATH)
+    if not os.path.isfile(TOKENIZER_PATH):
+        datamodule = torch.load(DATAMODULE_PATH)
+        datamodule.train_tokenizer()
+
 
 def parse_args():
     parser = argparse.ArgumentParser(allow_abbrev=True, formatter_class=argparse.RawTextHelpFormatter)
 
-    parser.add_argument("-m", "-max-epochs", help="limit the number of training epochs", type=int, dest="max_epochs")
-    parser.add_argument("-g", "-gpus", metavar="GPUS", type=int, choices=list(range(torch.cuda.device_count())),
-                        help="ids of gpus to train on, if not provided, then trains on cpu", nargs="+", dest="gpus")
-    parser.add_argument("-l", "-log", help="whether to save logs of run to w&b logger, default False", default=False,
+    parser.add_argument("gpus", type=int, default=None,
+                        help=f"Ids of gpus in range 0..{torch.cuda.device_count()} to train on, "
+                             "if not provided, then trains on cpu", nargs="*")
+    parser.add_argument("-l", "-log", help="Whether to save logs of run to w&b logger, default False", default=False,
                         action="store_true", dest="log")
-    parser.add_argument("-width", help="width of images, default 1024", default=1024, type=int)
-    parser.add_argument("-height", help="height of images, default 128", default=128, type=int)
-    parser.add_argument("-r", "-randomize", default=5, type=int, dest="random_magnitude", choices=range(10),
-                        help="add random augments to images of provided magnitude in range 0..9, default 5")
-    parser.add_argument("-b", "-batch-size", help="batch size, default 16", default=16,
-                        type=int, dest="batch_size")
+    parser.add_argument("-m", "-max-epochs", help="Limit the number of training epochs", type=int, dest="max_epochs")
+
+    datamodule_args = ["image_width", "image_height", "batch_size", "random_magnitude"]
+
+    datamodule = torch.load(DATAMODULE_PATH)
+    parser.add_argument("-d", metavar="X", nargs=4, dest="datamodule_args", type=int,
+                        help="Create new datamodule and exit, current parameters:\n" +
+                             "\n".join(f"{arg}\t{datamodule.hparams[arg]}" for arg in datamodule_args))
+
     transformer_args = [("num_encoder_layers", 6), ("num_decoder_layers", 6), ("d_model", 512), ("nhead", 8),
                         ("dim_feedforward", 2048), ("dropout", 0.1)]
-    parser.add_argument("-t", "-transformer-args", dest="transformer_args", nargs='+', default=[],
-                        help="transformer init args:\n" + "\n".join(f"{k}\t{v}" for k, v in transformer_args))
+    parser.add_argument("-t", metavar="X", dest="transformer_args", nargs=len(transformer_args),
+                        help="Transformer init args, reference values:\n" +
+                             "\n".join(f"{k}\t{v}" for k, v in transformer_args))
 
     args = parser.parse_args()
-    for i, parameter in enumerate(args.transformer_args):
-        transformer_args[i][1] = parameter
-    args.transformer_args = dict(transformer_args)
+
+    if args.datamodule_args:
+        args.datamodule_args = dict(zip(datamodule_args, args.datamodule_args))
+
+    if args.transformer_args:
+        args.transformer_args = dict(zip(list(zip(*transformer_args))[0], args.transformer_args))
+    else:
+        args.transformer_args = dict(transformer_args)
+
     return args
 
 
 def main():
+    check_setup()
     args = parse_args()
-    datamodule = LatexImageDataModule(image_width=args.width, image_height=args.height,
-                                      batch_size=args.batch_size, random_magnitude=args.random_magnitude)
-    datamodule.prepare_data()
+    if args.datamodule_args:
+        datamodule = LatexImageDataModule(image_width=args.datamodule_args["image_width"],
+                                          image_height=args.datamodule_args["image_height"],
+                                          batch_size=args.datamodule_args["batch_size"],
+                                          random_magnitude=args.datamodule_args["random_magnitude"])
+        datamodule.train_tokenizer()
+        tex_tokenizer = torch.load(TOKENIZER_PATH)
+        print(f"Vocabulary size {tex_tokenizer.get_vocab_size()}")
+        torch.save(datamodule, DATAMODULE_PATH)
+        return
 
+    datamodule = torch.load(DATAMODULE_PATH)
+    tex_tokenizer = torch.load(TOKENIZER_PATH)
+    logger = None
+    callbacks = []
     if args.log:
         logger = WandbLogger(f"img2tex", log_model=True)
-        callbacks = [LogImageTexCallback(logger),
-                     LearningRateMonitor(logging_interval='step'),
+        callbacks = [LogImageTexCallback(logger, top_k=10, max_length=20),
+                     LearningRateMonitor(logging_interval="step"),
                      ModelCheckpoint(save_top_k=10,
                                      monitor="val_loss",
                                      mode="min",
                                      filename="img2tex-{epoch:02d}-{val_loss:.2f}")]
-    else:
-        logger = None
-        callbacks = []
 
     trainer = Trainer(max_epochs=args.max_epochs,
                       accelerator="cpu" if args.gpus is None else "gpu",
                       gpus=args.gpus,
                       logger=logger,
-                      strategy="ddp",
+                      strategy="ddp_find_unused_parameters_false",
                       enable_progress_bar=True,
                       default_root_dir=TRAINER_DIR,
                       callbacks=callbacks,
                       check_val_every_n_epoch=5)
 
-    transformer = Transformer(num_encoder_layers=args.transformer_args['num_encoder_layers'],
-                              num_decoder_layers=args.transformer_args['num_decoder_layers'],
-                              d_model=args.transformer_args['d_model'],
-                              nhead=args.transformer_args['nhead'],
-                              dim_feedforward=args.transformer_args['dim_feedforward'],
-                              dropout=args.transformer_args['dropout'],
-                              image_width=datamodule.hparams['image_width'],
-                              image_height=datamodule.hparams['image_height'],
-                              tgt_vocab_size=datamodule.tex_tokenizer.get_vocab_size(),
-                              pad_idx=datamodule.tex_tokenizer.token_to_id("[PAD]"))
+    transformer = Transformer(num_encoder_layers=args.transformer_args["num_encoder_layers"],
+                              num_decoder_layers=args.transformer_args["num_decoder_layers"],
+                              d_model=args.transformer_args["d_model"],
+                              nhead=args.transformer_args["nhead"],
+                              dim_feedforward=args.transformer_args["dim_feedforward"],
+                              dropout=args.transformer_args["dropout"],
+                              image_width=datamodule.hparams["image_width"],
+                              image_height=datamodule.hparams["image_height"],
+                              tgt_vocab_size=tex_tokenizer.get_vocab_size(),
+                              pad_idx=tex_tokenizer.token_to_id("[PAD]"))
 
     trainer.fit(transformer, datamodule=datamodule)
-    trainer.save_checkpoint(os.path.join(TRAINER_DIR, "best_model.ckpt"))
+    trainer.test(transformer, datamodule=datamodule)
 
 
 if __name__ == "__main__":

utils.py

@@ -1,15 +1,19 @@
 from constants import TOKENIZER_PATH
+from data_preprocessing import RandomizeImageTransform
 
 import einops
 import random
 from pytorch_lightning.callbacks import Callback
 import torch
+import torch.nn.functional as F
 from torchvision import transforms
 
 
 class LogImageTexCallback(Callback):
-    def __init__(self, logger):
+    def __init__(self, logger, top_k, max_length):
         self.logger = logger
+        self.top_k = top_k
+        self.max_length = max_length
         self.tex_tokenizer = torch.load(TOKENIZER_PATH)
         self.tensor_to_PIL = transforms.ToPILImage()
 
@@ -18,101 +22,64 @@ class LogImageTexCallback(Callback):
             return
         sample_id = random.randint(0, len(batch['images']) - 1)
         image = batch['images'][sample_id]
-        tex_predicted, tex_ids = decode(transformer, self.tex_tokenizer, image)
+        texs_predicted, texs_ids = beam_search_decode(transformer, image, transform_image=False, top_k=self.top_k,
+                                                      max_length=self.max_length)
         image = self.tensor_to_PIL(image)
-        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}"])
+        tex_true = self.tex_tokenizer.decode(list(batch['tex_ids'][sample_id].to('cpu', torch.int)))
+        self.logger.log_image(key="samples", images=[image],
+                              caption=[f"True: {tex_true}\nPredicted: " + "\n".join(texs_predicted)])
 
 
-# parser.add_argument(
-#     "-t", "-tune", help="whether to tune model for batch size before training, default False", default=False,
-#     action="store_true", dest="tune"
-# )
+@torch.inference_mode()
+def beam_search_decode(transformer, image, transform_image=True, top_k=10, max_length=100):
+    """Performs decoding maintaining k best candidates"""
+    assert torch.is_tensor(image) and len(image.shape) == 3, "Image must be a 3 dimensional tensor (c h w)"
 
-# 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)
-# TUNER_DIR = "resources/pl_tuner_checkpoints"
-# from pytorch_lightning import  seed_everything
-#     parser.add_argument(
-#         "-d", "-deterministic", help="whether to seed all rngs for reproducibility, default False", default=False,
-#         action="store_true", dest="deterministic"
-#     )
-#     if args.deterministic:
-#         seed_everything(42, workers=True)
-# def generate_normalize_transform(dataset: TexImageDataset):
-#     """Returns a normalize layer with mean and std computed after iterating over dataset"""
-#
-#     mean = 0
-#     std = 0
-#     for item in tqdm.tqdm(dataset, "Computing dataset image stats"):
-#         image = item['image']
-#         mean += image.mean()
-#         std += image.std()
-#
-#     mean /= len(dataset)
-#     std /= len(dataset)
-#     normalize = T.Normalize(mean, std)
-#     return normalize
-# 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)
+    def get_tgt_padding_mask(tgt):
+        mask = tgt == tex_tokenizer.token_to_id("[SEP]")
+        mask = torch.cumsum(mask, dim=1)
+        mask = mask.to(transformer.device, torch.bool)
+        return mask
 
+    src = einops.rearrange(image, "c h w -> () c h w").to(transformer.device)
+    if transform_image:
+        image_transform = RandomizeImageTransform(width=transformer.hparams["image_width"],
+                                                  height=transformer.hparams["image_width"],
+                                                  random_magnitude=0)
+        src = image_transform(src)
+    memory = transformer.encode(src)
 
-@torch.inference_mode()
-def decode(transformer, image):
     tex_tokenizer = torch.load(TOKENIZER_PATH)
-    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
+    candidates_tex_ids = [[tex_tokenizer.token_to_id("[CLS]")]]
+    candidates_log_prob = torch.tensor([0], dtype=torch.float, device=transformer.device)
+
+    while candidates_tex_ids[0][-1] != tex_tokenizer.token_to_id("[SEP]") and len(candidates_tex_ids[0]) < max_length:
+        candidates_tex_ids = torch.tensor(candidates_tex_ids, dtype=torch.float, device=transformer.device)
+        tgt_mask = transformer.transformer.generate_square_subsequent_mask(candidates_tex_ids.shape[1]).to(
+            transformer.device, torch.bool)
+        shared_memories = einops.repeat(memory, f"one n d_model -> ({candidates_tex_ids.shape[0]} one) n d_model")
+        outs = transformer.decode(tgt=candidates_tex_ids,
+                                  memory=shared_memories,
+                                  tgt_mask=tgt_mask,
+                                  memory_mask=None,
+                                  tgt_padding_mask=get_tgt_padding_mask(candidates_tex_ids))
+        outs = einops.rearrange(outs, 'b n prob -> b prob n')[:, :, -1]
+        vocab_size = outs.shape[1]
+        outs = F.log_softmax(outs, dim=1)
+        outs += einops.rearrange(candidates_log_prob, "prob -> prob ()")
+        outs = einops.rearrange(outs, 'b prob -> (b prob)')
+        candidates_log_prob, indices = torch.topk(outs, k=top_k)
+
+        new_candidates = []
+        for index in indices:
+            candidate_id, token_id = divmod(index.item(), vocab_size)
+            new_candidates.append(candidates_tex_ids[candidate_id].to(int).tolist() + [token_id])
+        candidates_tex_ids = new_candidates
+
+    candidates_tex_ids = torch.tensor(candidates_tex_ids)
+    padding_mask = get_tgt_padding_mask(candidates_tex_ids).cpu()
+    candidates_tex_ids = candidates_tex_ids.masked_fill(
+        padding_mask & (candidates_tex_ids != tex_tokenizer.token_to_id("[SEP]")),
+        tex_tokenizer.token_to_id("[PAD]")).tolist()
+    texs = tex_tokenizer.decode_batch(candidates_tex_ids, skip_special_tokens=True)
+    return texs, candidates_tex_ids