moved constants to separate file, organized tokenizer

constants.py

@@ -0,0 +1,11 @@
+PDFLATEX = "/external2/dkkoshman/venv/texlive/2022/bin/x86_64-linux/pdflatex"
+GHOSTSCRIPT = "/external2/dkkoshman/venv/local/gs/bin/gs"
+
+DATA_DIR = "data"
+LATEX_PATH = "resources/latex.json"
+TRAINER_DIR = "resources/trainer"
+TOKENIZER_PATH = "resources/tokenizer.pt"
+
+NUM_DATALOADER_WORKERS = 4
+PERSISTENT_WORKERS = True  # whether to shut down workers at the end of epoch
+PIN_MEMORY = False  # probably causes cuda oom error if True

data_generator.py

@@ -1,3 +1,5 @@
+from constants import DATA_DIR, LATEX_PATH, PDFLATEX, GHOSTSCRIPT
+
 import json
 from multiprocessing import Pool
 import os
@@ -7,11 +9,6 @@ import subprocess
 import random
 import tqdm
 
-DATA_DIR = "data"
-LATEX_PATH = "resources/latex.json"
-PDFLATEX = "/external2/dkkoshman/venv/texlive/2022/bin/x86_64-linux/pdflatex"
-GHOSTSCRIPT = "/external2/dkkoshman/venv/local/gs/bin/gs"
-
 
 def generate_equation(latex, size, max_depth):
     """

data_preprocessing.py

@@ -1,4 +1,4 @@
-from data_generator import DATA_DIR
+from constants import DATA_DIR, TOKENIZER_PATH, NUM_DATALOADER_WORKERS, PERSISTENT_WORKERS, PIN_MEMORY
 
 import einops
 import os
@@ -9,14 +9,8 @@ import torchvision
 import torchvision.transforms as T
 from torch.utils.data import Dataset, DataLoader
 import tqdm
-import random
 import re
 
-TOKENIZER_PATH = "resources/tokenizer.pt"
-NUM_WORKERS = 4
-PERSISTENT_WORKERS = True  # whether to shut down workers at the end of epoch
-PIN_MEMORY = False  # probably causes cuda oom error if True
-
 
 class TexImageDataset(Dataset):
     """Image and tex dataset."""
@@ -147,18 +141,12 @@ def generate_tex_tokenizer(dataloader):
 class LatexImageDataModule(pl.LightningDataModule):
     def __init__(self, image_width, image_height, batch_size, random_magnitude):
         super().__init__()
-        image_transform = RandomizeImageTransform(image_width, image_height, random_magnitude)
-        tex_transform = ExtractEquationFromTexTransform()
-
-        self.train_dataset = TexImageDataset(DATA_DIR, image_transform, tex_transform)
-        self.val_dataset = TexImageDataset(DATA_DIR, image_transform, tex_transform)
-        self.test_dataset = TexImageDataset(DATA_DIR, image_transform, tex_transform)
-
-        train_indices, val_indices, test_indices = self.train_val_test_split(len(self.train_dataset))
-        self.train_dataset = torch.utils.data.Subset(self.train_dataset, train_indices)
-        self.val_dataset = torch.utils.data.Subset(self.val_dataset, val_indices)
-        self.test_dataset = torch.utils.data.Subset(self.test_dataset, test_indices)
 
+        dataset = TexImageDataset(root_dir=DATA_DIR,
+                                  image_transform=RandomizeImageTransform(image_width, image_height, random_magnitude),
+                                  tex_transform=ExtractEquationFromTexTransform())
+        self.train_dataset, self.val_dataset, self.test_dataset = torch.utils.data.random_split(
+            dataset, [len(dataset) * 18 // 20, len(dataset) // 20, len(dataset) // 20])
         self.batch_size = batch_size
         self.save_hyperparameters()
 
@@ -167,27 +155,17 @@ class LatexImageDataModule(pl.LightningDataModule):
         print(f"Vocabulary size: {tokenizer.get_vocab_size()}")
         torch.save(tokenizer, TOKENIZER_PATH)
 
-    def setup(self, stage=None):
-        self.tex_tokenizer = torch.load(TOKENIZER_PATH)
-        self.collate_fn = BatchCollator(self.tex_tokenizer)
-
-    @staticmethod
-    def train_val_test_split(size, train_fraction=.8, val_fraction=.1):
-        indices = list(range(size))
-        random.shuffle(indices)
-        train_split = int(size * train_fraction)
-        val_split = train_split + int(size * val_fraction)
-        return indices[:train_split], indices[train_split: val_split], indices[val_split:]
+    def _shared_dataloader(self, dataset, **kwargs):
+        tex_tokenizer = torch.load(TOKENIZER_PATH)
+        collate_fn = BatchCollator(tex_tokenizer)
+        return DataLoader(dataset, batch_size=self.batch_size, collate_fn=collate_fn, pin_memory=PIN_MEMORY,
+                          num_workers=NUM_DATALOADER_WORKERS, persistent_workers=PERSISTENT_WORKERS, **kwargs)
 
     def train_dataloader(self):
-        return DataLoader(self.train_dataset, batch_size=self.batch_size, collate_fn=self.collate_fn,
-                          pin_memory=PIN_MEMORY, num_workers=NUM_WORKERS, persistent_workers=PERSISTENT_WORKERS,
-                          shuffle=True)
+        return self._shared_dataloader(self.train_dataset, shuffle=True)
 
     def val_dataloader(self):
-        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)
+        return self._shared_dataloader(self.val_dataset)
 
     def test_dataloader(self):
-        return DataLoader(self.test_dataset, batch_size=self.batch_size, collate_fn=self.collate_fn,
-                          pin_memory=PIN_MEMORY, num_workers=NUM_WORKERS, persistent_workers=PERSISTENT_WORKERS)
+        return self._shared_dataloader(self.test_dataset)

model.py

@@ -30,7 +30,6 @@ class AddPositionalEncoding(nn.Module):
     def forward(self, batch):
         seq_len = batch.size(1)
         positional_encodings = self.positional_encodings[:seq_len, :]
-        # implicit batch broadcasting
         return batch + positional_encodings
 
 
@@ -125,7 +124,6 @@ class Transformer(pl.LightningModule):
         self.src_tok_emb = ImageEmbedding(d_model, image_width, image_height, patch_size=16, dropout=dropout)
         self.tgt_tok_emb = TexEmbedding(d_model, tgt_vocab_size, dropout=dropout)
         self.generator = nn.Linear(d_model, 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()

train.py

@@ -1,16 +1,15 @@
+from constants import TRAINER_DIR
 from data_preprocessing import LatexImageDataModule
 from model import Transformer
 from utils import LogImageTexCallback
 
 import argparse
 import os
-from pytorch_lightning.callbacks import LearningRateMonitor
+from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint
 from pytorch_lightning.loggers import WandbLogger
 from pytorch_lightning import Trainer
 import torch
 
-TRAINER_DIR = "resources/pl_trainer_checkpoints"
-
 
 # TODO: update python, maybe model doesnt train bc of ignore special index in CrossEntropyLoss?
 # crop image, adjust brightness, make tex tokens always decodable,
@@ -47,11 +46,15 @@ def main():
     datamodule = LatexImageDataModule(image_width=args.width, image_height=args.height,
                                       batch_size=args.batch_size, random_magnitude=args.random_magnitude)
     datamodule.prepare_data()
-    datamodule.setup()
+
     if args.log:
         logger = WandbLogger(f"img2tex", log_model=True)
-        callbacks = [LogImageTexCallback(logger, datamodule.tex_tokenizer),
-                     LearningRateMonitor(logging_interval='step')]
+        callbacks = [LogImageTexCallback(logger),
+                     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 = []
@@ -63,7 +66,8 @@ def main():
                       strategy="ddp",
                       enable_progress_bar=True,
                       default_root_dir=TRAINER_DIR,
-                      callbacks=callbacks)
+                      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'],
@@ -77,7 +81,6 @@ def main():
                               pad_idx=datamodule.tex_tokenizer.token_to_id("[PAD]"))
 
     trainer.fit(transformer, datamodule=datamodule)
-    trainer.test(datamodule=datamodule, ckpt_path='best')
     trainer.save_checkpoint(os.path.join(TRAINER_DIR, "best_model.ckpt"))
 
 

utils.py

@@ -1,4 +1,4 @@
-from model import Transformer
+from constants import TOKENIZER_PATH
 
 import einops
 import random
@@ -8,9 +8,9 @@ from torchvision import transforms
 
 
 class LogImageTexCallback(Callback):
-    def __init__(self, logger, tex_tokenizer):
+    def __init__(self, logger):
         self.logger = logger
-        self.tex_tokenizer = tex_tokenizer
+        self.tex_tokenizer = torch.load(TOKENIZER_PATH)
         self.tensor_to_PIL = transforms.ToPILImage()
 
     def on_validation_batch_start(self, trainer, transformer, batch, batch_idx, dataloader_idx):
@@ -102,7 +102,8 @@ class LogImageTexCallback(Callback):
 
 
 @torch.inference_mode()
-def decode(transformer, tex_tokenizer, image):
+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: