working backend

data_generator.py

@@ -1,5 +1,3 @@
-from train import DATA_DIR, LATEX_PATH
-
 import json
 from multiprocessing import Pool
 import os
@@ -9,6 +7,9 @@ import subprocess
 import random
 import tqdm
 
+DATA_DIR = 'data'
+LATEX_PATH = 'resources/latex.json'
+
 
 class DotDict(dict):
     """dot.notation access to dictionary attributes"""
@@ -168,7 +169,7 @@ def generate_data(examples_count) -> None:
     :examples_count: - how many latex - image examples to generate
     """
 
-    filenames = set(f"{i:0{len(str(examples_count - 1))}d}" for i in range(examples_count)),
+    filenames = set(f"{i:0{len(str(examples_count - 1))}d}" for i in range(examples_count))
     directory = os.path.abspath(DATA_DIR)
     latex_path = os.path.abspath(LATEX_PATH)
     with open(latex_path) as file:

data_preprocessing.py

@@ -1,4 +1,4 @@
-from train import DATASET_PATH, DATA_DIR, BATCH_SIZE, TEX_VOCAB_SIZE
+from data_generator import DATA_DIR
 
 import einops
 import os
@@ -9,9 +9,14 @@ import torchvision
 import torchvision.transforms as T
 from torch.utils.data import Dataset, DataLoader
 import tqdm
-from typing import Optional
+import random
 import re
 
+TEX_VOCAB_SIZE = 300
+BATCH_SIZE = 16
+IMAGE_WIDTH = 1024
+IMAGE_HEIGHT = 128
+
 
 class TexImageDataset(Dataset):
     """Image and tex dataset."""
@@ -89,7 +94,7 @@ class BatchCollator(object):
 class StandardizeImageTransform(object):
     """Pad and crop image to a given size, grayscale and invert"""
 
-    def __init__(self, width=1024, height=128):
+    def __init__(self, width=IMAGE_WIDTH, height=IMAGE_HEIGHT):
         self.standardize = T.Compose((
             T.Resize(height),
             T.Grayscale(),
@@ -106,7 +111,7 @@ class StandardizeImageTransform(object):
 class RandomizeImageTransform(object):
     """Standardize image and randomly augment"""
 
-    def __init__(self, width=1024, height=128, random_magnitude=5):
+    def __init__(self, width=IMAGE_WIDTH, height=IMAGE_HEIGHT, random_magnitude=5):
         self.transform = T.Compose((
             T.ColorJitter(brightness=random_magnitude / 10),
             T.Resize(height),
@@ -138,10 +143,10 @@ class ExtractEquationFromTexTransform(object):
         return equation
 
 
-def generate_tex_tokenizer(dataset: TexImageDataset, vocab_size=300):
+def generate_tex_tokenizer(dataset, vocab_size):
     """Returns a tokenizer trained on texs from given dataset"""
 
-    texs = list(tqdm.tqdm((item['tex'] for item in dataset), "Training tokenizer"))
+    texs = list(tqdm.tqdm((item['tex'] for item in dataset), "Training tokenizer", total=len(dataset)))
 
     os.environ['TOKENIZERS_PARALLELISM'] = 'false'
     tokenizer = tokenizers.Tokenizer(tokenizers.models.BPE(unk_token="[UNK]"))
@@ -164,31 +169,49 @@ def generate_tex_tokenizer(dataset: TexImageDataset, vocab_size=300):
 
 
 class LatexImageDataModule(pl.LightningDataModule):
-    def prepare_data(self) -> None:
-        # download or something
-        ...
-
-    def setup(self, stage: Optional[str] = None) -> None:
-        tex_transform = ExtractEquationFromTexTransform()
-        dataset = TexImageDataset(DATA_DIR, tex_transform=tex_transform)
+    def __init__(self):
+        super().__init__()
+        torch.manual_seed(0)
 
-        self.train_dataset, self.val_dataset, self.test_dataset = torch.utils.data.random_split(
-            dataset,
-            [len(dataset) - 2 * len(dataset) // 10, len(dataset) // 10, len(dataset) // 10]
+        self.train_dataset = TexImageDataset(
+            root_dir=DATA_DIR,
+            image_transform=RandomizeImageTransform(),
+            tex_transform=ExtractEquationFromTexTransform()
         )
-        self.train_dataset.image_transform = RandomizeImageTransform()
-        self.val_dataset.image_transform = StandardizeImageTransform()
-        self.test_dataset.image_transform = StandardizeImageTransform()
-        # image_normalize = generate_normalize_transform(self.train_dataset), compose?
+        self.val_dataset = TexImageDataset(
+            root_dir=DATA_DIR,
+            image_transform=StandardizeImageTransform(),
+            tex_transform=ExtractEquationFromTexTransform()
+        )
+        self.test_dataset = TexImageDataset(
+            root_dir=DATA_DIR,
+            image_transform=StandardizeImageTransform(),
+            tex_transform=ExtractEquationFromTexTransform()
+        )
+        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)
 
         self.tex_tokenizer = generate_tex_tokenizer(self.train_dataset, vocab_size=TEX_VOCAB_SIZE)
         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 train_dataloader(self):
-        return DataLoader(self.train_dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=self.collate_fn)
+        return DataLoader(self.train_dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=self.collate_fn,
+                          num_workers=8, pin_memory=True, persistent_workers=False)
 
     def val_dataloader(self):
-        return DataLoader(self.val_dataset, batch_size=BATCH_SIZE, shuffle=False, collate_fn=self.collate_fn)
+        return DataLoader(self.val_dataset, batch_size=BATCH_SIZE, collate_fn=self.collate_fn, num_workers=8,
+                          pin_memory=True, persistent_workers=False)
 
     def test_dataloader(self):
-        return DataLoader(self.test_dataset, batch_size=BATCH_SIZE, shuffle=False, collate_fn=self.collate_fn)
+        return DataLoader(self.test_dataset, batch_size=BATCH_SIZE, collate_fn=self.collate_fn, num_workers=8,
+                          pin_memory=True, persistent_workers=False)

model.py

@@ -2,7 +2,6 @@ from einops.layers.torch import Rearrange
 import einops
 import math
 import pytorch_lightning as pl
-from pytorch_lightning.utilities.types import TRAIN_DATALOADERS
 import torch.nn as nn
 import torch
 
@@ -101,9 +100,6 @@ class ImageEncoder(nn.Module):
 
 
 class Transformer(pl.LightningModule):
-    def train_dataloader(self) -> TRAIN_DATALOADERS:
-        pass
-
     def __init__(self,
                  num_encoder_layers: int,
                  num_decoder_layers: int,
@@ -139,6 +135,20 @@ class Transformer(pl.LightningModule):
                                 src_padding_mask, tgt_padding_mask)
         return self.generator(outs)
 
+    def general_step(self, batch):
+        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)
+        src_padding_mask = None
+        tgt_padding_mask = batch['tex_attention_masks'][:, :-1]
+        outs = self(src, tgt_input, src_mask, tgt_mask, src_padding_mask, tgt_padding_mask)
+        loss = self.loss_fn(einops.rearrange(outs, 'b n prob -> b prob n'), tgt_output.long())
+        return loss
+
     def training_step(self, batch, batch_idx):
         src = batch['images']
         tgt = batch['tex_ids']
@@ -154,5 +164,16 @@ class Transformer(pl.LightningModule):
         self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)
         return loss
 
+    def validation_step(self, batch, batch_idx):
+        loss = self.general_step(batch)
+        self.log("val_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)
+        return loss
+
+    def test_step(self, batch, batch_idx):
+        loss = self.general_step(batch)
+        self.log("test_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)
+        return loss
+
     def configure_optimizers(self):
+        # TODO write scheduler
         return torch.optim.Adam(self.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9)

train.py

@@ -1,5 +1,5 @@
 from data_generator import generate_data
-from data_preprocessing import LatexImageDataModule
+from data_preprocessing import LatexImageDataModule, IMAGE_WIDTH, IMAGE_HEIGHT
 from model import Transformer
 
 import argparse
@@ -7,31 +7,26 @@ import pytorch_lightning as pl
 from pytorch_lightning.loggers import WandbLogger
 import torch
 
-DATA_DIR = 'data'
-LATEX_PATH = 'resources/latex.json'
-DATASET_PATH = 'resources/dataset'
-IMAGE_WIDTH = 1024
-IMAGE_HEIGHT = 128
-TEX_VOCAB_SIZE = 300
-BATCH_SIZE = 16
+DATASET_PATH = 'resources/dataset.pt'
 
 
 def main():
-    torch.manual_seed(0)
-
     parser = argparse.ArgumentParser("Trainer")
-    parser.add_argument("-generate-new", help="number of new files to generate", type=int)
+    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", help="list of gpu ids to train on", type=int, nargs='+', dest="gpus",
+                        choices=list(range(torch.cuda.device_count())), default=[0])
     args = parser.parse_args()
 
-    if args.generate_new is not None:
-        generate_data(args.generate_new)
+    if args.new_dataset is not None:
+        generate_data(args.new_dataset)
         datamodule = LatexImageDataModule()
         torch.save(datamodule, DATASET_PATH)
     else:
         datamodule = torch.load(DATASET_PATH)
 
     wandb_logger = WandbLogger()
-    trainer = pl.Trainer(max_epochs=2, accelerator='gpu', gpus=1, logger=wandb_logger)
+    trainer = pl.Trainer(max_epochs=2, accelerator='gpu', gpus=args.gpus, logger=wandb_logger, strategy='ddp_spawn')
     transformer = Transformer(
         num_encoder_layers=3,
         num_decoder_layers=3,