dedicated generate.py script

data_generator.py

@@ -7,109 +7,76 @@ import subprocess
 import random
 import tqdm
 
-DATA_DIR = 'data'
-LATEX_PATH = 'resources/latex.json'
-
-
-class DotDict(dict):
-    """dot.notation access to dictionary attributes"""
-    __getattr__ = dict.get
-    __setattr__ = dict.__setitem__
-    __delattr__ = dict.__delitem__
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        if len(args) > 0 and isinstance(args[0], dict):
-            for key, value in self.items():
-                if isinstance(value, dict):
-                    self.__setitem__(key, DotDict(value))
-
-
-def _generate_equation(size_left, depth_left, latex, tokens):
-    if size_left <= 0:
-        return ""
-
-    equation = ""
-    pairs, scopes, special = latex.pairs, latex.scopes, latex.special
-    weights = [3, depth_left > 0, depth_left > 0]
-    group, = random.choices([tokens, pairs, scopes], weights=weights)
-
-    if group is tokens:
-        equation += ' '.join([
-            random.choice(tokens),
-            _generate_equation(size_left - 1, depth_left, latex, tokens)
-        ])
-        return equation
-
-    post_scope_size = round(abs(random.gauss(0, size_left / 2)))
-    size_left -= post_scope_size + 1
-
-    if group is pairs:
-        pair = random.choice(pairs)
-        equation += ' '.join([
-            pair[0],
-            _generate_equation(size_left, depth_left - 1, latex, tokens),
-            pair[1],
-            _generate_equation(post_scope_size, depth_left, latex, tokens)
-        ])
-        return equation
+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"
 
-    elif group is scopes:
-        scope_type, scope_group = random.choice(list(scopes.items()))
-        scope_operator = random.choice(scope_group)
-        equation += scope_operator
 
-        if scope_type == 'single':
-            equation += ' '.join([
-                special.left_bracket,
-                _generate_equation(size_left, depth_left - 1, latex, tokens)
+def generate_equation(latex, size, max_depth):
+    """
+        Generates a random latex equation
+        -------
+        params:
+        :latex: -- dict with tokens to generate equation from
+        :size: -- approximate size of equation
+        :max_depth: -- max brackets and scope depth
+        """
+
+    tokens, pairs, scopes = latex["tokens"], latex["pairs"], latex["scope_manipulators"]
+
+    def _generate_equation_recursive(size_left=size, depth_used=0):
+        if size_left <= 0:
+            return ""
+
+        equation = ""
+        group, = random.choices([tokens, pairs, scopes],
+                                weights=[max_depth + 1, max_depth > depth_used, max_depth > depth_used])
+
+        if group is tokens:
+            equation += " ".join([
+                random.choice(tokens),
+                _generate_equation_recursive(size_left - 1, depth_used)
             ])
-
-        elif scope_type == 'double_no_delimiters':
-            equation += ' '.join([
-                special.left_bracket,
-                _generate_equation(size_left // 2, depth_left - 1, latex, tokens),
-                special.right_bracket + special.left_bracket,
-                _generate_equation(size_left // 2, depth_left - 1, latex, tokens)
+            return equation
+
+        post_scope_size = round(abs(random.gauss(0, size_left / 2)))
+        size_left -= post_scope_size + 1
+
+        if group is pairs:
+            pair = random.choice(pairs)
+            equation += " ".join([
+                pair[0],
+                _generate_equation_recursive(size_left, depth_used + 1),
+                pair[1],
+                _generate_equation_recursive(post_scope_size, depth_used)
             ])
+            return equation
 
-        elif scope_type == 'double_with_delimiters':
-            equation += ' '.join([
-                special.caret,
-                special.left_bracket,
-                _generate_equation(size_left // 2, depth_left - 1, latex, tokens),
-                special.right_bracket,
-                special.underscore,
-                special.left_bracket,
-                _generate_equation(size_left // 2, depth_left - 1, latex, tokens)
-            ])
+        elif group is scopes:
+            scope_type, scope_group = random.choice(list(scopes.items()))
+            scope_operator = random.choice(scope_group)
+            equation += scope_operator
 
-        equation += ' '.join([
-            special.right_bracket,
-            _generate_equation(post_scope_size, depth_left, latex, tokens)
-        ])
-        return equation
+            if scope_type == "single":
+                equation += "{ " + _generate_equation_recursive(size_left, depth_used + 1)
 
+            elif scope_type == "double_no_delimiters":
+                equation += "{ " + _generate_equation_recursive(size_left // 2, depth_used + 1) + " } { " + \
+                            _generate_equation_recursive(size_left // 2, depth_used + 1)
 
-def generate_equation(latex: DotDict, size, depth=3):
-    """
-    Generates a random latex equation
-    -------
-    params:
-    :latex: -- dict with tokens to generate equation from
-    :size: -- approximate size of equation
-    :depth: -- max brackets and scope depth
-    """
-    tokens = [token for group in ['chars', 'greek', 'functions', 'operators', 'spaces']
-              for token in latex[group]]
-    equation = _generate_equation(size, depth, latex, tokens)
-    return equation
+            elif scope_type == "double_with_delimiters":
+                equation += "^ { " + _generate_equation_recursive(size_left // 2, depth_used + 1) + " } _ { " + \
+                            _generate_equation_recursive(size_left // 2, depth_used + 1)
+
+            equation += _generate_equation_recursive(post_scope_size, depth_used) + " }"
 
+        return equation
 
-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"
-                   ):
+    return _generate_equation_recursive()
+
+
+def generate_image(directory, latex, filename, max_depth, equation_length, distribution_fraction):
     """
     Generates a random tex file and corresponding image
     -------
@@ -117,41 +84,47 @@ def generate_image(directory: str, latex: dict, filename: str, max_length=20, eq
     :directory: -- dir where to save files
     :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
+    :max_depth: -- max nested level of tex scopes
+    :equation_length: -- max length of equation
+    :distribution_fraction: -- fraction of whole available tex tokens to use
     """
+    fracture = lambda sequence: sequence[:max(1, int(len(sequence) * distribution_fraction))]
+    for group in ["tokens", "pairs", "fonts", "font_sizes"]:
+        latex[group] = fracture(latex[group])
+    for key, value in list(latex["scope_manipulators"].items()):
+        latex["scope_manipulators"]['key'] = fracture(value)
+
+    size = random.randint((equation_length + 1) // 2, equation_length)
+    equation = generate_equation(latex, size=size, max_depth=max_depth)
+
+    font, font_options = random.choice(latex["fonts"])
+    font_option = random.choice([""] + font_options)
+    font_size = random.choice(latex["font_sizes"])
+    template = string.Template(latex["template"])
+    tex = template.substitute(font=font, font_option=font_option, fontsize=font_size, equation=equation)
+
     filepath = os.path.join(directory, filename)
-    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)
-    fontsize = random.choice(latex.fontsizes)
-    equation = generate_equation(latex, equation_length, depth=equation_depth)
-    tex = template.substitute(font=font, font_option=font_option, fontsize=fontsize, equation=equation)
-
-    with open(f"{filepath}.tex", mode='w') as file:
+    with open(f"{filepath}.tex", mode="w") as file:
         file.write(tex)
 
     try:
         pdflatex_process = subprocess.run(
-            f"{pdflatex} -output-directory={directory} {filepath}.tex".split(),
+            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())
+        os.remove(filepath + ".tex")
         return
 
     if pdflatex_process.returncode != 0:
-        subprocess.run(f'rm {filepath}.tex'.split())
+        os.remove(filepath + ".tex")
         return
 
     subprocess.run(
-        f"{ghostscript} -sDEVICE=png16m -dTextAlphaBits=4 -r200 -dSAFER -dBATCH -dNOPAUSE -o {filepath}.png {filepath}.pdf".split(),
+        f"{GHOSTSCRIPT} -sDEVICE=png16m -dTextAlphaBits=4 -r200 -dSAFER -dBATCH -dNOPAUSE"
+        f" -o {filepath}.png {filepath}.pdf".split(),
         stderr=subprocess.DEVNULL,
         stdout=subprocess.DEVNULL,
     )
@@ -161,41 +134,33 @@ def _generate_image_wrapper(args):
     return generate_image(*args)
 
 
-def generate_data(examples_count) -> None:
+def generate_data(examples_count, max_depth, equation_length, distribution_fraction) -> None:
     """
     Clears a directory and generates a latex dataset in given directory
-    -------
-    params:
-    :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))
     directory = os.path.abspath(DATA_DIR)
-    latex_path = os.path.abspath(LATEX_PATH)
-    with open(latex_path) as file:
-        latex = json.load(file)
+    shutil.rmtree(DATA_DIR)
+    os.mkdir(DATA_DIR)
 
-    shutil.rmtree(directory)
-    os.mkdir(directory)
+    with open(LATEX_PATH) as file:
+        latex = json.load(file)
 
-    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()))
+    filenames = set(f"{i:0{len(str(examples_count - 1))}d}" for i in range(examples_count))
+    files_before = set(os.listdir())
 
-    files_before = _get_current_relevant_files()
     while filenames:
         with Pool() as pool:
             list(tqdm.tqdm(
-                pool.imap(_generate_image_wrapper, ((directory, latex, filename) for filename in sorted(filenames))),
+                pool.imap(_generate_image_wrapper,
+                          ((directory, latex, filename, max_depth, equation_length, distribution_fraction) 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)
+        filenames -= set(
+            os.path.splitext(filename)[0] for filename in os.listdir(directory) if filename.endswith(".png"))
+
+    for file in set(i.path for i in os.scandir(DATA_DIR)) | set(os.listdir()) - files_before:
+        if any(file.endswith(ext) for ext in [".aux", ".pdf", ".log", ".sh"]):
+            os.remove(file)

data_preprocessing.py

@@ -12,10 +12,7 @@ import tqdm
 import random
 import re
 
-TEX_VOCAB_SIZE = 300
-IMAGE_WIDTH = 1024
-IMAGE_HEIGHT = 128
-BATCH_SIZE = 16
+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
@@ -60,22 +57,6 @@ class TexImageDataset(Dataset):
         return {"image": image, "tex": tex}
 
 
-def generate_normalize_transform(dataset: TexImageDataset):
-    """Returns a normalize layer with mean and std computed after iterating over dataset"""
-
-    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):
     """Image, tex batch collator"""
 
@@ -94,39 +75,30 @@ class BatchCollator(object):
         return {'images': images, 'tex_ids': tex_ids, 'tex_attention_masks': attention_masks}
 
 
-class StandardizeImageTransform(object):
-    """Pad and crop image to a given size, grayscale and invert"""
-
-    def __init__(self, width=IMAGE_WIDTH, height=IMAGE_HEIGHT):
-        self.standardize = T.Compose((
-            T.Resize(height),
-            T.Grayscale(),
-            T.functional.invert,
-            T.CenterCrop((height, width)),
-            T.ConvertImageDtype(torch.float32)
-        ))
-
-    def __call__(self, image):
-        image = self.standardize(image)
-        return image
-
-
 class RandomizeImageTransform(object):
     """Standardize image and randomly augment"""
 
-    def __init__(self, width=IMAGE_WIDTH, height=IMAGE_HEIGHT, random_magnitude=5):
-        assert random_magnitude > 0
-        eps = 0.01
-        self.transform = T.Compose((
-            T.ColorJitter(brightness=((1 - eps) / (random_magnitude + eps), 1 - eps)),
-            T.Resize(height),
-            T.Grayscale(),
-            T.functional.invert,
-            T.CenterCrop((height, width)),
-            torch.Tensor.contiguous,
-            T.RandAugment(magnitude=random_magnitude),
-            T.ConvertImageDtype(torch.float32)
-        ))
+    def __init__(self, width, height, random_magnitude):
+        if random_magnitude > 0:
+            self.transform = T.Compose((
+                T.ColorJitter(brightness=random_magnitude / 10, contrast=random_magnitude / 10,
+                              saturation=random_magnitude / 10, hue=min(0.5, random_magnitude / 10)),
+                T.Resize(height),
+                T.Grayscale(),
+                T.functional.invert,
+                T.CenterCrop((height, width)),
+                torch.Tensor.contiguous,
+                T.RandAugment(magnitude=random_magnitude),
+                T.ConvertImageDtype(torch.float32)
+            ))
+        else:
+            self.transform = T.Compose((
+                T.Resize(height),
+                T.Grayscale(),
+                T.functional.invert,
+                T.CenterCrop((height, width)),
+                T.ConvertImageDtype(torch.float32)
+            ))
 
     def __call__(self, image):
         image = self.transform(image)
@@ -148,7 +120,7 @@ class ExtractEquationFromTexTransform(object):
         return equation
 
 
-def generate_tex_tokenizer(dataloader, vocab_size):
+def generate_tex_tokenizer(dataloader):
     """Returns a tokenizer trained on texs from given dataset"""
 
     texs = list(tqdm.tqdm((batch['tex'] for batch in dataloader), "Training tokenizer", total=len(dataloader)))
@@ -156,7 +128,6 @@ def generate_tex_tokenizer(dataloader, vocab_size):
     os.environ['TOKENIZERS_PARALLELISM'] = 'false'
     tokenizer = tokenizers.Tokenizer(tokenizers.models.BPE(unk_token="[UNK]"))
     tokenizer_trainer = tokenizers.trainers.BpeTrainer(
-        vocab_size=vocab_size,
         special_tokens=["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]"]
     )
     tokenizer.pre_tokenizer = tokenizers.pre_tokenizers.Whitespace()
@@ -174,34 +145,30 @@ def generate_tex_tokenizer(dataloader, vocab_size):
 
 
 class LatexImageDataModule(pl.LightningDataModule):
-    def __init__(self, batch_size=BATCH_SIZE):
+    def __init__(self, image_width, image_height, batch_size, random_magnitude):
         super().__init__()
-        torch.manual_seed(0)
-        self.batch_size = batch_size
+        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)
 
-        self.train_dataset = TexImageDataset(
-            root_dir=DATA_DIR,
-            image_transform=RandomizeImageTransform(),
-            tex_transform=ExtractEquationFromTexTransform()
-        )
-        self.val_dataset = TexImageDataset(
-            root_dir=DATA_DIR,
-            image_transform=RandomizeImageTransform(),
-            tex_transform=ExtractEquationFromTexTransform()
-        )
-        self.test_dataset = TexImageDataset(
-            root_dir=DATA_DIR,
-            image_transform=RandomizeImageTransform(),
-            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(
-            DataLoader(self.train_dataset, batch_size=32, num_workers=16),
-            vocab_size=TEX_VOCAB_SIZE)
+        self.batch_size = batch_size
+        self.save_hyperparameters()
+
+    def prepare_data(self):
+        tokenizer = generate_tex_tokenizer(DataLoader(self.train_dataset, batch_size=32, num_workers=16))
+        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
@@ -213,8 +180,9 @@ 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=self.batch_size, shuffle=True, collate_fn=self.collate_fn,
-                          pin_memory=PIN_MEMORY, num_workers=NUM_WORKERS, persistent_workers=PERSISTENT_WORKERS)
+        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)
 
     def val_dataloader(self):
         return DataLoader(self.val_dataset, batch_size=self.batch_size, collate_fn=self.collate_fn,

generate.py

@@ -0,0 +1,23 @@
+from data_generator import generate_data
+
+import argparse
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Clear old dataset and generate new one")
+    parser.add_argument("size", help="size of new dataset", type=int)
+    parser.add_argument("depth", help="max_depth scope depth of generated equation, no less than 1", type=int)
+    parser.add_argument("length", help="length of equation will be in range length/2..length", type=int)
+    parser.add_argument("fraction", help="fraction of tex vocab to sample tokens from, float in range 0..1", type=float)
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+    generate_data(examples_count=args.size, max_depth=args.depth, equation_length=args.length,
+                  distribution_fraction=args.fraction)
+
+
+if __name__ == "__main__":
+    main()

model.py

@@ -11,10 +11,6 @@ class AddPositionalEncoding(nn.Module):
     def __init__(self, d_model, max_sequence_len=5000):
         super().__init__()
 
-        # pos - position in sequence, i - index of element embedding
-        # PE(pos, 2i) = sin(pos / 10000**(2i / d_model)) = sin(pos * e**(2i * (-log(10000))/d_model))
-        # PE(pos, 2i+1) = cos(pos / 10000**(2i / d_model)) = cos(pos * e**(2i * (-log(10000))/d_model))
-
         positions = torch.arange(max_sequence_len)
         even_embedding_indices = torch.arange(0, d_model, 2)
 
@@ -103,7 +99,7 @@ class Transformer(pl.LightningModule):
     def __init__(self,
                  num_encoder_layers: int,
                  num_decoder_layers: int,
-                 emb_size: int,
+                 d_model: int,
                  nhead: int,
                  image_width: int,
                  image_height: int,
@@ -114,7 +110,7 @@ class Transformer(pl.LightningModule):
                  ):
         super().__init__()
 
-        self.transformer = nn.Transformer(d_model=emb_size,
+        self.transformer = nn.Transformer(d_model=d_model,
                                           nhead=nhead,
                                           num_encoder_layers=num_encoder_layers,
                                           num_decoder_layers=num_decoder_layers,
@@ -125,10 +121,10 @@ class Transformer(pl.LightningModule):
             if p.dim() > 1:
                 nn.init.xavier_uniform_(p)
 
-        self.d_model = emb_size
-        self.src_tok_emb = ImageEmbedding(emb_size, image_width, image_height, patch_size=16, dropout=dropout)
-        self.tgt_tok_emb = TexEmbedding(emb_size, tgt_vocab_size, dropout=dropout)
-        self.generator = nn.Linear(emb_size, tgt_vocab_size)
+        self.d_model = d_model
+        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)

train.py

@@ -1,73 +1,57 @@
-from data_generator import generate_data
-from data_preprocessing import LatexImageDataModule, IMAGE_WIDTH, IMAGE_HEIGHT
+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.loggers import WandbLogger
-from pytorch_lightning import Trainer, seed_everything
+from pytorch_lightning import Trainer
 import torch
 
-DATASET_PATH = "resources/dataset.pt"
 TRAINER_DIR = "resources/pl_trainer_checkpoints"
-TUNER_DIR = "resources/pl_tuner_checkpoints"
-BEST_MODEL_CHECKPOINT = "best_model.ckpt"
 
 
+# TODO: update python, maybe model doesnt train bc of ignore special index in CrossEntropyLoss?
+# crop image, adjust brightness, make tex tokens always decodable,
+# save only datamodule state?, ensemble last checkpoints, early stopping
+
 def parse_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "-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="ids of gpus to train on, if not provided then trains on cpu", nargs="+",
-        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"
-    )
-    # parser.add_argument(
-    #     "-t", "-tune", help="whether to tune model for batch size before training, default False", default=False,
-    #     action="store_true", dest="tune"
-    # )
+    parser = argparse.ArgumentParser(allow_abbrev=True, formatter_class=argparse.RawTextHelpFormatter)
+
+    parser.add_argument("-m", "-max-epochs", help="limit the number of training epochs", type=int, dest="max_epochs")
+    parser.add_argument("-g", "-gpus", metavar="GPUS", type=int, choices=list(range(torch.cuda.device_count())),
+                        help="ids of gpus to train on, if not provided, then trains on cpu", nargs="+", dest="gpus")
+    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("-width", help="width of images, default 1024", default=1024, type=int)
+    parser.add_argument("-height", help="height of images, default 128", default=128, type=int)
+    parser.add_argument("-r", "-randomize", default=5, type=int, dest="random_magnitude", choices=range(10),
+                        help="add random augments to images of provided magnitude in range 0..9, default 5")
+    parser.add_argument("-b", "-batch-size", help="batch size, default 16", default=16,
+                        type=int, dest="batch_size")
+    transformer_args = [("num_encoder_layers", 6), ("num_decoder_layers", 6), ("d_model", 512), ("nhead", 8),
+                        ("dim_feedforward", 2048), ("dropout", 0.1)]
+    parser.add_argument("-t", "-transformer-args", dest="transformer_args", nargs='+', default=[],
+                        help="transformer init args:\n" + "\n".join(f"{k}\t{v}" for k, v in transformer_args))
 
     args = parser.parse_args()
+    for i, parameter in enumerate(args.transformer_args):
+        transformer_args[i][1] = parameter
+    args.transformer_args = dict(transformer_args)
     return args
 
 
-# TODO: update python, maybe model doesnt train bc of ignore special index in CrossEntropyLoss?
-# crop image, adjust brightness, lr warmup?, make tex tokens always decodable,
-# take loss that doesn't punish so much for offsets, take a look at weights,
-
 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)
-        datamodule = LatexImageDataModule()
-        torch.save(datamodule, DATASET_PATH)
-    else:
-        datamodule = torch.load(DATASET_PATH)
-
+    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, datamodule.tex_tokenizer),
+                     LearningRateMonitor(logging_interval='step')]
     else:
         logger = None
         callbacks = []
@@ -79,24 +63,22 @@ def main():
                       strategy="ddp",
                       enable_progress_bar=True,
                       default_root_dir=TRAINER_DIR,
-                      callbacks=callbacks,
-                      )
-
-    transformer = Transformer(num_encoder_layers=3,
-                              num_decoder_layers=3,
-                              emb_size=512,
-                              nhead=8,
-                              image_width=IMAGE_WIDTH,
-                              image_height=IMAGE_HEIGHT,
+                      callbacks=callbacks)
+
+    transformer = Transformer(num_encoder_layers=args.transformer_args['num_encoder_layers'],
+                              num_decoder_layers=args.transformer_args['num_decoder_layers'],
+                              d_model=args.transformer_args['d_model'],
+                              nhead=args.transformer_args['nhead'],
+                              dim_feedforward=args.transformer_args['dim_feedforward'],
+                              dropout=args.transformer_args['dropout'],
+                              image_width=datamodule.hparams['image_width'],
+                              image_height=datamodule.hparams['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,
-                              )
+                              pad_idx=datamodule.tex_tokenizer.token_to_id("[PAD]"))
 
     trainer.fit(transformer, datamodule=datamodule)
-    trainer.test(datamodule=datamodule)
-    trainer.save_checkpoint(BEST_MODEL_CHECKPOINT)
+    trainer.test(datamodule=datamodule, ckpt_path='best')
+    trainer.save_checkpoint(os.path.join(TRAINER_DIR, "best_model.ckpt"))
 
 
 if __name__ == "__main__":

utils.py

@@ -22,57 +22,84 @@ class LogImageTexCallback(Callback):
         image = self.tensor_to_PIL(image)
         tex_true = self.tex_tokenizer.decode(list(batch['tex_ids'][sample_id].to('cpu', torch.int)),
                                              skip_special_tokens=True)
-        self.logger.log_image(key="samples", images=[image],
-                              caption=[f"True: {tex_true}\nPredicted: {tex_predicted}\nIds: {tex_ids}"])
+        self.logger.log_image(key="samples", images=[image], caption=[f"True: {tex_true}\nPredicted: {tex_predicted}"])
 
-# if args.new_dataset:
-    #     datamodule.batch_size = 1
-    #     transformer_for_tuning = TransformerTuner(**transformer.hparams).cuda()
-    #     tuner = Trainer(accelerator="gpu" if args.gpus else "cpu",
-    #                     gpus=args.gpus,
-    #                     strategy=TRAINER_STRATEGY,
-    #                     enable_progress_bar=True,
-    #                     enable_checkpointing=False,
-    #                     auto_scale_batch_size=True,
-    #                     num_sanity_val_steps=0,
-    #                     logger=False
-    #                     )
-    #     tuner.tune(transformer_for_tuning, datamodule=datamodule)
-    #     torch.save(datamodule, DATASET_PATH)
-
-class _TransformerTuner(Transformer):
-    """
-    When using trainer.tune, batches from dataloader get passed directly to forward,
-    so this subclass takes care of that
-    """
 
-    def forward(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)
-        memory_mask = None
-        src_padding_mask = None
-        tgt_padding_mask = batch['tex_attention_masks'][:, :-1]
-        tgt_padding_mask = tgt_padding_mask.masked_fill(
-            tgt_padding_mask == 0, float('-inf')
-        ).masked_fill(
-            tgt_padding_mask == 1, 0
-        )
+# parser.add_argument(
+#     "-t", "-tune", help="whether to tune model for batch size before training, default False", default=False,
+#     action="store_true", dest="tune"
+# )
 
-        src = self.src_tok_emb(src)
-        tgt_input = self.tgt_tok_emb(tgt_input)
-        outs = self.transformer(src, tgt_input, src_mask, tgt_mask, memory_mask, src_padding_mask, tgt_padding_mask)
-        outs = self.generator(outs)
-
-        loss = self.loss_fn(einops.rearrange(outs, 'b n prob -> b prob n'), tgt_output.long())
-        return loss
+# if args.new_dataset:
+#     datamodule.batch_size = 1
+#     transformer_for_tuning = TransformerTuner(**transformer.hparams).cuda()
+#     tuner = Trainer(accelerator="gpu" if args.gpus else "cpu",
+#                     gpus=args.gpus,
+#                     strategy=TRAINER_STRATEGY,
+#                     enable_progress_bar=True,
+#                     enable_checkpointing=False,
+#                     auto_scale_batch_size=True,
+#                     num_sanity_val_steps=0,
+#                     logger=False
+#                     )
+#     tuner.tune(transformer_for_tuning, datamodule=datamodule)
+#     torch.save(datamodule, DATASET_PATH)
+# TUNER_DIR = "resources/pl_tuner_checkpoints"
+# from pytorch_lightning import  seed_everything
+#     parser.add_argument(
+#         "-d", "-deterministic", help="whether to seed all rngs for reproducibility, default False", default=False,
+#         action="store_true", dest="deterministic"
+#     )
+#     if args.deterministic:
+#         seed_everything(42, workers=True)
+# def generate_normalize_transform(dataset: TexImageDataset):
+#     """Returns a normalize layer with mean and std computed after iterating over dataset"""
+#
+#     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 _TransformerTuner(Transformer):
+#     """
+#     When using trainer.tune, batches from dataloader get passed directly to forward,
+#     so this subclass takes care of that
+#     """
+#
+#     def forward(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)
+#         memory_mask = None
+#         src_padding_mask = None
+#         tgt_padding_mask = batch['tex_attention_masks'][:, :-1]
+#         tgt_padding_mask = tgt_padding_mask.masked_fill(
+#             tgt_padding_mask == 0, float('-inf')
+#         ).masked_fill(
+#             tgt_padding_mask == 1, 0
+#         )
+#
+#         src = self.src_tok_emb(src)
+#         tgt_input = self.tgt_tok_emb(tgt_input)
+#         outs = self.transformer(src, tgt_input, src_mask, tgt_mask, memory_mask, src_padding_mask, tgt_padding_mask)
+#         outs = self.generator(outs)
+#
+#         loss = self.loss_fn(einops.rearrange(outs, 'b n prob -> b prob n'), tgt_output.long())
+#         return loss
+#
+#     def validation_step(self, batch, batch_idx):
+#         return self(batch, batch_idx)
 
-    def validation_step(self, batch, batch_idx):
-        return self(batch, batch_idx)
 
 @torch.inference_mode()
 def decode(transformer, tex_tokenizer, image):
@@ -87,4 +114,4 @@ def decode(transformer, tex_tokenizer, image):
         next_id = outs[0, :, -1].argmax().item()
         tex_ids.append(next_id)
     tex = tex_tokenizer.decode(tex_ids, skip_special_tokens=True)
-    return tex, tex_ids
+    return tex, tex_ids