lr logger

data_preprocessing.py

@@ -15,7 +15,7 @@ import re
 TEX_VOCAB_SIZE = 300
 IMAGE_WIDTH = 1024
 IMAGE_HEIGHT = 128
-BATCH_SIZE = 8
+BATCH_SIZE = 16
 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
@@ -146,10 +146,10 @@ class ExtractEquationFromTexTransform(object):
         return equation
 
 
-def generate_tex_tokenizer(dataset, vocab_size):
+def generate_tex_tokenizer(dataloader, vocab_size):
     """Returns a tokenizer trained on texs from given dataset"""
 
-    texs = list(tqdm.tqdm((item['tex'] for item in dataset), "Training tokenizer", total=len(dataset)))
+    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]"))
@@ -197,7 +197,9 @@ class LatexImageDataModule(pl.LightningDataModule):
         self.val_dataset = torch.utils.data.Subset(self.val_dataset, val_indices)
         self.test_dataset = torch.utils.data.Subset(self.test_dataset, test_indices)
 
-        self.tex_tokenizer = generate_tex_tokenizer(self.train_dataset, vocab_size=TEX_VOCAB_SIZE)
+        self.tex_tokenizer = generate_tex_tokenizer(
+            DataLoader(self.train_dataset, batch_size=32, num_workers=16),
+            vocab_size=TEX_VOCAB_SIZE)
         self.collate_fn = BatchCollator(self.tex_tokenizer)
 
     @staticmethod

model.py

@@ -111,8 +111,7 @@ class Transformer(pl.LightningModule):
                  pad_idx: int,
                  dim_feedforward: int = 512,
                  dropout: float = .1,
-                 learning_rate=1e-3,
-                 tex_tokenizer=None
+                 learning_rate: float = 1e-3
                  ):
         super().__init__()
 
@@ -133,7 +132,6 @@ class Transformer(pl.LightningModule):
         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=None, tgt_mask=None, memory_mask=None, src_padding_mask=None,
                 tgt_padding_mask=None):
@@ -185,6 +183,10 @@ class Transformer(pl.LightningModule):
         scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=50, T_mult=1)
         return [optimizer], [scheduler]
 
+    # def configure_optimizers(self):
+    #     optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
+    #     return optimizer
+
 
 class _TransformerTuner(Transformer):
     """

train.py

@@ -4,10 +4,10 @@ from model import Transformer, _TransformerTuner
 from utils import LogImageTexCallback
 
 import argparse
+from pytorch_lightning.callbacks import LearningRateMonitor
 from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger
 from pytorch_lightning import Trainer, seed_everything
 import torch
-import wandb
 
 DATASET_PATH = "resources/dataset.pt"
 TRAINER_DIR = "resources/pl_trainer_checkpoints"
@@ -58,13 +58,15 @@ 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)
-        callbacks = [LogImageTexCallback(logger, datamodule.tex_tokenizer)]
+        callbacks = [
+            LogImageTexCallback(logger, datamodule.tex_tokenizer),
+            LearningRateMonitor(logging_interval='step')
+        ]
     else:
         logger = None
         callbacks = []
@@ -88,15 +90,10 @@ def main():
                               tgt_vocab_size=datamodule.tex_tokenizer.get_vocab_size(),
                               pad_idx=datamodule.tex_tokenizer.token_to_id("[PAD]"),
                               dim_feedforward=512,
-                              dropout=0.1
+                              dropout=0.1,
+                              learning_rate=1e-3
                               )
 
-    # 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()

utils.py

@@ -18,4 +18,4 @@ class LogImageTexCallback(Callback):
         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}"])
+        self.logger.log_image(key="samples", images=[image], caption=[f"True: {tex_true}\n Predicted: {tex_predicted}"])