big changes: transformer, pytorch lightning, argparser

data_generator.py

@@ -1,10 +1,13 @@
+from train import DATA_DIR, LATEX_PATH
+
 import json
 from multiprocessing import Pool
 import os
+import shutil
 import string
 import subprocess
 import random
-from typing import Iterable
+import tqdm
 
 
 class DotDict(dict):
@@ -102,7 +105,7 @@ def generate_equation(latex: DotDict, size, depth=3):
     return equation
 
 
-def generate_image(directory: str, latex: DotDict, filename: str, max_length=20, equation_depth=3,
+def generate_image(directory: str, latex: dict, filename: str, max_length=20, equation_depth=3,
                    pdflatex: str = "/external2/dkkoshman/venv/texlive/2022/bin/x86_64-linux/pdflatex",
                    ghostscript: str = "/external2/dkkoshman/venv/local/gs/bin/gs"
                    ):
@@ -111,18 +114,16 @@ def generate_image(directory: str, latex: DotDict, filename: str, max_length=20,
     -------
     params:
     :directory: -- dir where to save files
-    :latex: -- DotDict with parameters to generate tex
+    :latex: -- dict with parameters to generate tex
     :filename: -- absolute filename for the generated files
     :max_length: -- max size of equation
     :equation_depth: -- max nested level of tex scopes
     :pdflatex: -- path to pdflatex
     :ghostscript: -- path to ghostscript
     """
-    # TODO ARGPARSE
     filepath = os.path.join(directory, filename)
-
-    equation_length = random.randint(1, max_length)
-
+    equation_length = random.randint(max_length // 2, max_length)
+    latex = DotDict(latex)
     template = string.Template(latex.template)
     font, font_options = random.choice(latex.fonts)
     font_option = random.choice([''] + font_options)
@@ -130,70 +131,70 @@ def generate_image(directory: str, latex: DotDict, filename: str, max_length=20,
     equation = generate_equation(latex, equation_length, depth=equation_depth)
     tex = template.substitute(font=font, font_option=font_option, fontsize=fontsize, equation=equation)
 
-    files_before = set(os.listdir(directory))
     with open(f"{filepath}.tex", mode='w') as file:
         file.write(tex)
 
-    pr1 = subprocess.run(
-        f"{pdflatex} -output-directory={directory} {filepath}.tex".split(),
-        stderr=subprocess.PIPE,
-    )
+    try:
+        pdflatex_process = subprocess.run(
+            f"{pdflatex} -output-directory={directory} {filepath}.tex".split(),
+            stderr=subprocess.DEVNULL,
+            stdout=subprocess.DEVNULL,
+            timeout=1
+        )
+    except subprocess.TimeoutExpired:
+        subprocess.run(f'rm {filepath}.tex'.split())
+        return
 
-    files_after = set(os.listdir(directory))
-    if pr1.returncode != 0:
-        files_to_delete = files_after - files_before
-        if files_to_delete:
-            subprocess.run(['rm'] + [os.path.join(directory, file) for file in files_to_delete])
-        print(pr1.stderr.decode(), tex)
+    if pdflatex_process.returncode != 0:
+        subprocess.run(f'rm {filepath}.tex'.split())
         return
 
-    pr2 = subprocess.run(
+    subprocess.run(
         f"{ghostscript} -sDEVICE=png16m -dTextAlphaBits=4 -r200 -dSAFER -dBATCH -dNOPAUSE -o {filepath}.png {filepath}.pdf".split(),
-        stderr=subprocess.PIPE,
+        stderr=subprocess.DEVNULL,
+        stdout=subprocess.DEVNULL,
     )
 
-    files_to_delete = files_after - files_before - {filename + '.png', filename + '.tex'}
-    if files_to_delete:
-        subprocess.run(['rm'] + [os.path.join(directory, file) for file in files_to_delete])
-    assert (pr2.returncode == 0)
 
+def _generate_image_wrapper(args):
+    return generate_image(*args)
 
-def generate_data(
-        filenames: Iterable[str],
-        directory: str,
-        latex_path: str,
-        overwrite: bool = False
-) -> None:
+
+def generate_data(examples_count) -> None:
     """
-    Generates a latex dataset in given directory
+    Clears a directory and generates a latex dataset in given directory
     -------
     params:
-    :filenames: - iterable of filenames to create, without extension
-    :directory: - where to create
-    :latex_path: - full path to latex json
-    :overwrite: - whether to overwrite existing files
+    :examples_count: - how many latex - image examples to generate
     """
-    subprocess.run(". /external2/dkkoshman/venv/bin/activate")
-    if not os.path.isabs(directory):
-        directory = os.path.join(os.getcwd(), directory)
-    if not os.path.isabs(latex_path):
-        latex_path = os.path.join(os.getcwd(), latex_path)
-
-    filenames = set(filenames)
-    if not overwrite:
-        existing = set(
-            filename for file in os.listdir(directory) for filename, ext in os.path.splitext(file) if ext == '.png'
-        )
-        filenames -= existing
 
+    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:
         latex = json.load(file)
-    latex = DotDict(latex)
 
+    shutil.rmtree(directory)
+    os.mkdir(directory)
+
+    def _get_current_relevant_files():
+        return set(os.path.join(directory, file) for file in os.listdir(directory)) | set(
+            os.path.abspath(file) for file in os.listdir(os.getcwd()))
+
+    files_before = _get_current_relevant_files()
     while filenames:
-        for name in filenames:
-            generate_image(directory, latex, name)
-        # with Pool() as pool:
-        #     pool.starmap(generate_image, ((directory, latex, name) for name in filenames))
-        existing = set(file.split('.')[0] for file in os.listdir(directory) if file.endswith('.png'))
+        with Pool() as pool:
+            list(tqdm.tqdm(
+                pool.imap(_generate_image_wrapper, ((directory, latex, filename) for filename in sorted(filenames))),
+                "Generating images",
+                total=len(filenames)
+            ))
+        existing = set(os.path.splitext(filename)[0] for filename in os.listdir(directory) if filename.endswith('.png'))
         filenames -= existing
+
+    files_after = _get_current_relevant_files()
+    files_to_delete = files_after - files_before
+    files_to_delete = list(os.path.join(directory, file) for file in files_to_delete if
+                           not file.endswith('.png') and not file.endswith('.tex'))
+    if files_to_delete:
+        subprocess.run(['rm'] + files_to_delete)

data_preprocessing.py

@@ -1,11 +1,15 @@
+from train import DATASET_PATH, DATA_DIR, BATCH_SIZE, TEX_VOCAB_SIZE
+
 import einops
 import os
+import pytorch_lightning as pl
 import tokenizers
 import torch
 import torchvision
 import torchvision.transforms as T
-from torch.utils.data import Dataset
+from torch.utils.data import Dataset, DataLoader
 import tqdm
+from typing import Optional
 import re
 
 
@@ -23,19 +27,10 @@ class TexImageDataset(Dataset):
         torch.multiprocessing.set_sharing_strategy('file_system')
         self.root_dir = root_dir
         self.filenames = sorted(set(
-            filename for file in os.listdir(root_dir) for filename, ext in os.path.splitext(file) if ext == '.png'
+            os.path.splitext(filename)[0] for filename in os.listdir(root_dir) if filename.endswith('.png')
         ))
         self.image_transform = image_transform
         self.tex_transform = tex_transform
-        self.tex_tokenizer = None
-        self.texs = []
-        for filename in tqdm.tqdm(self.filenames, "Preloading tex files"):
-            tex_path = os.path.join(self.root_dir, filename + '.tex')
-            with open(tex_path) as file:
-                tex = file.read()
-            if self.tex_transform:
-                tex = self.tex_transform(tex)
-            self.texs.append(tex)
 
     def __len__(self):
         return len(self.filenames)
@@ -43,29 +38,34 @@ class TexImageDataset(Dataset):
     def __getitem__(self, idx):
         filename = self.filenames[idx]
         image_path = os.path.join(self.root_dir, filename + '.png')
+        tex_path = os.path.join(self.root_dir, filename + '.tex')
+
+        with open(tex_path) as file:
+            tex = file.read()
+        if self.tex_transform:
+            tex = self.tex_transform(tex)
+
         image = torchvision.io.read_image(image_path)
         if self.image_transform:
             image = self.image_transform(image)
-        tex = self.texs[idx]
+
         return {"image": image, "tex": tex}
 
-    def subjoin_image_normalize_transform(self):
-        """Appends a normalize layer with mean and std computed after iterating over dataset"""
-        mean = 0
-        std = 0
-        for item in tqdm.tqdm(self):
-            image = item['image']
-            mean += image.mean()
-            std += image.std()
 
-        mean /= len(self)
-        std /= len(self)
-        normalize = T.Normalize(mean, std)
+def generate_normalize_transform(dataset: TexImageDataset):
+    """Returns a normalize layer with mean and std computed after iterating over dataset"""
 
-        if self.image_transform:
-            self.image_transform = T.Compose((self.image_transform, normalize))
-        else:
-            self.image_transform = normalize
+    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 BatchCollator(object):
@@ -138,10 +138,12 @@ class ExtractEquationFromTexTransform(object):
         return equation
 
 
-def generate_tex_tokenizer(texs, vocab_size=300):
-    """Returns a tokenizer trained on given tex strings"""
+def generate_tex_tokenizer(dataset: TexImageDataset, vocab_size=300):
+    """Returns a tokenizer trained on texs from given dataset"""
 
-    # os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+    texs = list(tqdm.tqdm((item['tex'] for item in dataset), "Training tokenizer"))
+
+    os.environ['TOKENIZERS_PARALLELISM'] = 'false'
     tokenizer = tokenizers.Tokenizer(tokenizers.models.BPE(unk_token="[UNK]"))
     tokenizer_trainer = tokenizers.trainers.BpeTrainer(
         vocab_size=vocab_size,
@@ -150,9 +152,43 @@ def generate_tex_tokenizer(texs, vocab_size=300):
     tokenizer.pre_tokenizer = tokenizers.pre_tokenizers.Whitespace()
     tokenizer.train_from_iterator(texs, trainer=tokenizer_trainer)
     tokenizer.post_processor = tokenizers.processors.TemplateProcessing(
-        single="$A [SEP]",
-        special_tokens=[("[SEP]", tokenizer.token_to_id("[SEP]"))]
+        single="[CLS] $A [SEP]",
+        special_tokens=[
+            ("[CLS]", tokenizer.token_to_id("[CLS]")),
+            ("[SEP]", tokenizer.token_to_id("[SEP]")),
+        ]
     )
     tokenizer.enable_padding(pad_id=tokenizer.token_to_id("[PAD]"), pad_token="[PAD]")
 
     return tokenizer
+
+
+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)
+
+        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.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.tex_tokenizer = generate_tex_tokenizer(self.train_dataset, vocab_size=TEX_VOCAB_SIZE)
+        self.collate_fn = BatchCollator(self.tex_tokenizer)
+
+    def train_dataloader(self):
+        return DataLoader(self.train_dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=self.collate_fn)
+
+    def val_dataloader(self):
+        return DataLoader(self.val_dataset, batch_size=BATCH_SIZE, shuffle=False, collate_fn=self.collate_fn)
+
+    def test_dataloader(self):
+        return DataLoader(self.test_dataset, batch_size=BATCH_SIZE, shuffle=False, collate_fn=self.collate_fn)

model.py

@@ -1,6 +1,8 @@
 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
 
@@ -62,6 +64,21 @@ class ImageEmbedding(nn.Module):
         return image_batch
 
 
+class TexEmbedding(nn.Module):
+    def __init__(self, d_model: int, vocab_size: int, dropout: float = .1):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, d_model)
+        self.add_positional_encoding = AddPositionalEncoding(d_model)
+        self.dropout = nn.Dropout(p=dropout)
+        self.d_model = d_model
+
+    def forward(self, tex_ids_batch):
+        tex_ids_batch = self.embedding(tex_ids_batch.long()) * math.sqrt(self.d_model)
+        tex_ids_batch = self.add_positional_encoding(tex_ids_batch)
+        tex_ids_batch = self.dropout(tex_ids_batch)
+        return tex_ids_batch
+
+
 class ImageEncoder(nn.Module):
     """
     Given an image, returns its vector representation.
@@ -83,7 +100,10 @@ class ImageEncoder(nn.Module):
         return self.encode(batch)
 
 
-class Seq2SeqTransformer(nn.Module):
+class Transformer(pl.LightningModule):
+    def train_dataloader(self) -> TRAIN_DATALOADERS:
+        pass
+
     def __init__(self,
                  num_encoder_layers: int,
                  num_decoder_layers: int,
@@ -92,39 +112,47 @@ class Seq2SeqTransformer(nn.Module):
                  image_width: int,
                  image_height: int,
                  tgt_vocab_size: int,
+                 pad_idx: int,
                  dim_feedforward: int = 512,
-                 dropout: float = 0.1):
-        super(Seq2SeqTransformer, self).__init__()
+                 dropout: float = .1):
+        super().__init__()
         self.transformer = nn.Transformer(d_model=emb_size,
                                           nhead=nhead,
                                           num_encoder_layers=num_encoder_layers,
                                           num_decoder_layers=num_decoder_layers,
                                           dim_feedforward=dim_feedforward,
-                                          dropout=dropout)
-        # TODO: share weights between generator and embedding
+                                          dropout=dropout,
+                                          batch_first=True)
+        for p in self.transformer.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
         self.generator = nn.Linear(emb_size, tgt_vocab_size)
         self.src_tok_emb = ImageEmbedding(emb_size, image_width, image_height, dropout=dropout)
-        self.tgt_tok_emb = TokenEmbedding(tgt_vocab_size, emb_size)
-
-    def forward(self,
-                src: Tensor,
-                trg: Tensor,
-                src_mask: Tensor,
-                tgt_mask: Tensor,
-                src_padding_mask: Tensor,
-                tgt_padding_mask: Tensor,
-                memory_key_padding_mask: Tensor):
-        src_emb = self.positional_encoding(self.src_tok_emb(src))
-        tgt_emb = self.positional_encoding(self.tgt_tok_emb(trg))
+        self.tgt_tok_emb = TexEmbedding(emb_size, tgt_vocab_size, dropout=dropout)
+        self.loss_fn = torch.nn.CrossEntropyLoss(ignore_index=pad_idx)
+
+    def forward(self, src, tgt, src_mask, tgt_mask, src_padding_mask, tgt_padding_mask):
+        src_emb = self.src_tok_emb(src)
+        tgt_emb = self.tgt_tok_emb(tgt)
         outs = self.transformer(src_emb, tgt_emb, src_mask, tgt_mask, None,
-                                src_padding_mask, tgt_padding_mask, memory_key_padding_mask)
+                                src_padding_mask, tgt_padding_mask)
         return self.generator(outs)
 
-    def encode(self, src: Tensor, src_mask: Tensor):
-        return self.transformer.encoder(self.positional_encoding(
-            self.src_tok_emb(src)), src_mask)
-
-    def decode(self, tgt: Tensor, memory: Tensor, tgt_mask: Tensor):
-        return self.transformer.decoder(self.positional_encoding(
-            self.tgt_tok_emb(tgt)), memory,
-            tgt_mask)
+    def training_step(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)
+        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())
+        self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)
+        return loss
+
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9)

train.py

@@ -1,32 +1,54 @@
 from data_generator import generate_data
-from data_preprocessing import TexImageDataset, RandomizeImageTransform, ExtractEquationFromTexTransform, \
-    BatchCollator, generate_tex_tokenizer
+from data_preprocessing import LatexImageDataModule
+from model import Transformer
 
+import argparse
+import pytorch_lightning as pl
+from pytorch_lightning.loggers import WandbLogger
 import torch
-from torch.utils.data import DataLoader
 
 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
 
-if __name__ == '__main__':
-    generate_data(
-        filenames=map(str, range(1000)),
-        directory=DATA_DIR,
-        latex_path=LATEX_PATH,
-    )
 
-    image_transform = RandomizeImageTransform()
-    tex_transform = ExtractEquationFromTexTransform()
-    dataset = TexImageDataset(DATA_DIR, image_transform=image_transform, tex_transform=tex_transform)
-    dataset.subjoin_image_normalize_transform()
-    train_dataset, test_dataset = torch.utils.data.random_split(
-        dataset,
-        [len(dataset) * 9 // 10, len(dataset) // 10]
+def main():
+    torch.manual_seed(0)
+
+    parser = argparse.ArgumentParser("Trainer")
+    parser.add_argument("-generate-new", help="number of new files to generate", type=int)
+    args = parser.parse_args()
+
+    if args.generate_new is not None:
+        generate_data(args.generate_new)
+        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)
+    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
     )
-    tex_tokenizer = generate_tex_tokenizer(dataset.texs)
-    collate_fn = BatchCollator(tex_tokenizer)
 
-    train_dataloader = DataLoader(train_dataset, batch_size=16, shuffle=True, num_workers=16,
-                                  collate_fn=collate_fn)
-    batch = next(iter(train_dataloader))
-    print(batch['texs'])
+    trainer.fit(transformer, datamodule=datamodule)
+    trainer.validate(datamodule=datamodule)
+    trainer.test(datamodule=datamodule)
+
+
+if __name__ == '__main__':
+    main()