tuned cli, added tuner

data_preprocessing.py

@@ -15,8 +15,10 @@ import re
 TEX_VOCAB_SIZE = 300
 IMAGE_WIDTH = 1024
 IMAGE_HEIGHT = 128
-BATCH_SIZE = 16
-NUM_WORKERS = 0
+BATCH_SIZE = 8
+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):
@@ -170,9 +172,10 @@ def generate_tex_tokenizer(dataset, vocab_size):
 
 
 class LatexImageDataModule(pl.LightningDataModule):
-    def __init__(self):
+    def __init__(self, batch_size=BATCH_SIZE):
         super().__init__()
         torch.manual_seed(0)
+        self.batch_size = batch_size
 
         self.train_dataset = TexImageDataset(
             root_dir=DATA_DIR,
@@ -206,16 +209,13 @@ class LatexImageDataModule(pl.LightningDataModule):
         return indices[:train_split], indices[train_split: val_split], indices[val_split:]
 
     def train_dataloader(self):
-        return DataLoader(self.train_dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=self.collate_fn,
-                          num_workers=NUM_WORKERS, )
-        # pin_memory=True, persistent_workers=True)
+        return DataLoader(self.train_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 val_dataloader(self):
-        return DataLoader(self.val_dataset, batch_size=BATCH_SIZE, collate_fn=self.collate_fn,
-                          num_workers=NUM_WORKERS, )
-        # pin_memory=True, persistent_workers=True)
+        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):
-        return DataLoader(self.test_dataset, batch_size=BATCH_SIZE, collate_fn=self.collate_fn,
-                          num_workers=NUM_WORKERS, )
-        # pin_memory=True, persistent_workers=True)
+        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)

model.py

@@ -110,8 +110,10 @@ class Transformer(pl.LightningModule):
                  tgt_vocab_size: int,
                  pad_idx: int,
                  dim_feedforward: int = 512,
-                 dropout: float = .1):
+                 dropout: float = .1,
+                 learning_rate=1e-4):
         super().__init__()
+
         self.transformer = nn.Transformer(d_model=emb_size,
                                           nhead=nhead,
                                           num_encoder_layers=num_encoder_layers,
@@ -127,6 +129,7 @@ class Transformer(pl.LightningModule):
         self.src_tok_emb = ImageEmbedding(emb_size, image_width, image_height, 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
 
     def forward(self, src, tgt, src_mask, tgt_mask, memory_mask, src_padding_mask, tgt_padding_mask):
         src = self.src_tok_emb(src)
@@ -174,4 +177,4 @@ class Transformer(pl.LightningModule):
 
     def configure_optimizers(self):
         # TODO write scheduler
-        return torch.optim.Adam(self.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9)
+        return torch.optim.Adam(self.parameters(), lr=self.learning_rate)

train.py

@@ -3,32 +3,44 @@ from data_preprocessing import LatexImageDataModule, IMAGE_WIDTH, IMAGE_HEIGHT
 from model import Transformer
 
 import argparse
-import pytorch_lightning as pl
 from pytorch_lightning.loggers import WandbLogger
+from pytorch_lightning import Trainer, seed_everything
 import torch
 
 DATASET_PATH = 'resources/dataset.pt'
 
 
-def main():
-    parser = argparse.ArgumentParser("Trainer", usage="run trainer")
+def parse_args():
+    parser = argparse.ArgumentParser()
     parser.add_argument(
-        "epochs", help="number of epochs to train", type=int
+        "-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", metavar="GPUS", help=f"list of gpu ids to train on in range 0..{torch.cuda.device_count()}",
-        type=int, nargs='+', dest="gpus", choices=list(range(torch.cuda.device_count())),
+        "-g", "-gpus", help=f"number of gpus to train on in range 0..{torch.cuda.device_count()}",
+        type=int, dest="gpus", choices=list(range(torch.cuda.device_count())),
     )
     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(
+        "-d", "-deterministic", help="whether to seed all rngs for reproducibility, default False", default=False,
+        action="store_true", dest="deterministic"
+    )
 
     args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+
+    if args.deterministic:
+        seed_everything(42, workers=True)
 
     if args.new_dataset is not None:
         generate_data(args.new_dataset)
@@ -39,25 +51,36 @@ 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_epochs{args.epochs}_size{len(datamodule)}_gpus{args.gpus}_v0") if args.log else None
-    trainer = pl.Trainer(max_epochs=args.epochs, accelerator='gpu' if args.gpus else 'cpu', gpus=args.gpus,
-                         logger=logger, strategy='ddp_spawn')
-    transformer = Transformer(
-        num_encoder_layers=3,
-        num_decoder_layers=3,
-        emb_size=512,
-        nhead=8,
-        image_width=IMAGE_WIDTH,
-        image_height=IMAGE_HEIGHT,
-        tgt_vocab_size=datamodule.tex_tokenizer.get_vocab_size(),
-        pad_idx=datamodule.tex_tokenizer.token_to_id("[PAD]"),
-        dim_feedforward=512,
-        dropout=0.1
-    )
 
+    logger = WandbLogger(f"img2tex", version='0') if args.log else False
+
+    trainer = Trainer(max_epochs=args.max_epochs,
+                      accelerator='gpu' if args.gpus else 'cpu',
+                      gpus=args.gpus,
+                      logger=logger,
+                      strategy='ddp',
+                      auto_scale_batch_size="power",
+                      auto_lr_find=True,
+                      auto_select_gpus=True,
+                      enable_progress_bar=True
+                      )
+
+    transformer = Transformer(num_encoder_layers=3,
+                              num_decoder_layers=3,
+                              emb_size=512,
+                              nhead=8,
+                              image_width=IMAGE_WIDTH,
+                              image_height=IMAGE_HEIGHT,
+                              tgt_vocab_size=datamodule.tex_tokenizer.get_vocab_size(),
+                              pad_idx=datamodule.tex_tokenizer.token_to_id("[PAD]"),
+                              dim_feedforward=512,
+                              dropout=0.1
+                              )
+
+    trainer.tune(transformer, datamodule=datamodule)
     trainer.fit(transformer, datamodule=datamodule)
-    # trainer.validate(datamodule=datamodule)
     trainer.test(datamodule=datamodule)
+    trainer.save_checkpoint("best_model.ckpt")
 
 
 if __name__ == '__main__':