app.py interface, made functions more independent, ensemble, working prototype

app.py

@@ -1,17 +1,29 @@
-import streamlit as st
+from constants import RESOURCES
+from data_preprocessing import RandomizeImageTransform
+from utils import beam_search_decode
 
-st.markdown("### Hello, world!")
-st.markdown("<img width=200px src='https://rozetked.me/images/uploads/dwoilp3BVjlE.jpg'>", unsafe_allow_html=True)
-# ^-- можно показывать пользователю текст, картинки, ограниченное подмножество html - всё как в jupyter
+import streamlit as st
+import PIL
+import torch
+import torchvision.transforms as T
 
-text = st.text_area("TEXT HERE")
-# ^-- показать текстовое поле. В поле text лежит строка, которая находится там в данный момент
+MODEL_PATH = RESOURCES + "/model_2tcuvfsj.pt"
 
-# from transformers import pipeline
-# pipe = pipeline("ner", "Davlan/distilbert-base-multilingual-cased-ner-hrl")
-# raw_predictions = pipe(text)
-# тут уже знакомый вам код с huggingface.transformers -- его можно заменить на что угодно от fairseq до catboost
+# TODO: make faster
+transformer = torch.load(MODEL_PATH)
+image_transform = T.Compose((
+    T.ToTensor(),
+    RandomizeImageTransform(width=transformer.hparams['image_width'],
+                            height=transformer.hparams['image_height'],
+                            random_magnitude=0)
+))
 
-# st.markdown(f"{raw_predictions}")
-st.markdown(f"Simon says {text}!")
-# выводим результаты модели в текстовое поле, на потеху пользователю
+st.markdown("### Image to TeX")
+st.image("resources/frontend/latex_example_1.png")
+file_png = st.file_uploader("Upload a PNG image", type=([".png"]))
+if file_png is not None:
+    image = PIL.Image.open(file_png)
+    image = image.convert("RGB")
+    tex = beam_search_decode(transformer, image, image_transform=image_transform)
+    st.latex(tex[0])
+    st.text(tex[0])

constants.py

@@ -1,11 +1,14 @@
 PDFLATEX = "/external2/dkkoshman/venv/texlive/2022/bin/x86_64-linux/pdflatex"
 GHOSTSCRIPT = "/external2/dkkoshman/venv/local/gs/bin/gs"
 
-DATA_DIR = "data"
-LATEX_PATH = "resources/latex.json"
-TRAINER_DIR = "resources/trainer"
-TOKENIZER_PATH = "resources/tokenizer.pt"
-DATAMODULE_PATH = "resources/datamodule.pt"
+DATA_DIR = "local/data"
+WANDB_DIR = "local/wandb"
+TRAINER_DIR = "local/trainer"
+
+RESOURCES = "resources"
+LATEX_PATH = RESOURCES + "/latex.json"
+TOKENIZER_PATH = RESOURCES + "/tokenizer.pt"
+DATAMODULE_PATH = RESOURCES + "/datamodule.pt"
 
 NUM_DATALOADER_WORKERS = 4
 PERSISTENT_WORKERS = True  # whether to shut down workers at the end of epoch

data_preprocessing.py

@@ -73,26 +73,17 @@ class RandomizeImageTransform(object):
     """Standardize image and randomly augment"""
 
     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)
-            ))
+        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, max_size=width),
+            T.Grayscale(),
+            T.functional.invert,
+            T.CenterCrop((height, width)),
+            torch.Tensor.contiguous,
+            T.RandAugment(magnitude=random_magnitude),
+            T.ConvertImageDtype(torch.float32)
+        ))
 
     def __call__(self, image):
         image = self.transform(image)

generate.py

@@ -1,23 +0,0 @@
-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

@@ -73,27 +73,6 @@ class TexEmbedding(nn.Module):
         return tex_ids_batch
 
 
-class ImageEncoder(nn.Module):
-    """
-    Given an image, returns its vector representation.
-    """
-
-    def __init__(self, image_width, image_height, d_model, num_layers=8):
-        super().__init__()
-        image_embedding = ImageEmbedding(d_model, image_width, image_height, patch_size=16, dropout=.1)
-        encoder_layer = nn.TransformerEncoderLayer(
-            d_model=d_model,
-            nhead=8,
-            dim_feedforward=2048,
-            batch_first=True
-        )
-        transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
-        self.encode = nn.Sequential(image_embedding, transformer_encoder)
-
-    def forward(self, batch):
-        return self.encode(batch)
-
-
 class Transformer(pl.LightningModule):
     def __init__(self,
                  num_encoder_layers: int,

train.py

@@ -1,7 +1,8 @@
-from constants import TRAINER_DIR, TOKENIZER_PATH, DATAMODULE_PATH
+from constants import TRAINER_DIR, TOKENIZER_PATH, DATAMODULE_PATH, WANDB_DIR, RESOURCES
+from data_generator import generate_data
 from data_preprocessing import LatexImageDataModule
 from model import Transformer
-from utils import LogImageTexCallback
+from utils import LogImageTexCallback, average_checkpoints
 
 import argparse
 import os
@@ -11,44 +12,60 @@ from pytorch_lightning import Trainer
 import torch
 
 
-# TODO: update python, make tex tokens always decodable, ensemble last checkpoints,
-#  clear checkpoint data build full dataset, train export model to torchscript write spaces interface
-
 def check_setup():
+    print(
+        "Disabling tokenizers parallelism because it can't be used before forking and I didn't bother to figure it out")
     os.environ["TOKENIZERS_PARALLELISM"] = "false"
     if not os.path.isfile(DATAMODULE_PATH):
+        print("Generating default datamodule")
         datamodule = LatexImageDataModule(image_width=1024, image_height=128, batch_size=16, random_magnitude=5)
         torch.save(datamodule, DATAMODULE_PATH)
     if not os.path.isfile(TOKENIZER_PATH):
+        print("Generating default tokenizer")
         datamodule = torch.load(DATAMODULE_PATH)
         datamodule.train_tokenizer()
 
 
 def parse_args():
-    parser = argparse.ArgumentParser(allow_abbrev=True, formatter_class=argparse.RawTextHelpFormatter)
-
-    parser.add_argument("gpus", type=int, default=None,
-                        help=f"Ids of gpus in range 0..{torch.cuda.device_count()} to train on, "
-                             "if not provided, then trains on cpu", nargs="*")
-    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("-m", "-max-epochs", help="Limit the number of training epochs", type=int, dest="max_epochs")
-
-    datamodule_args = ["image_width", "image_height", "batch_size", "random_magnitude"]
+    parser = argparse.ArgumentParser(description="Workflow: generate dataset, create datamodule, train model",
+                                     allow_abbrev=True, formatter_class=argparse.RawTextHelpFormatter)
+
+    parser.add_argument(
+        "gpus", type=int, help=f"Ids of gpus in range 0..{torch.cuda.device_count() - 1} to train on, "
+                               "if not provided,\nthen trains on cpu. To see current gpu load, run nvtop", nargs="*")
+    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(
+        "-m", "-max-epochs", help="Limit the number of training epochs", type=int, dest="max_epochs")
+
+    data_args = ["size", "depth", "length", "fraction"]
+    parser.add_argument(
+        "-n", metavar=tuple(map(str.upper, data_args)), nargs=4, dest="data_args",
+        type=lambda x: int(x) if x.isdigit() else float(x),
+        help="Clear old dataset, create new and exit, args:"
+             "\nsize\tsize of new dataset"
+             "\ndepth\tmax_depth scope depth of generated equation, no less than 1"
+             "\nlength\tlength of equation will be in range length/2..length"
+             "\nfraction\tfraction of tex vocab to sample tokens from, float in range 0..1")
 
     datamodule = torch.load(DATAMODULE_PATH)
-    parser.add_argument("-d", metavar="X", nargs=4, dest="datamodule_args", type=int,
-                        help="Create new datamodule and exit, current parameters:\n" +
-                             "\n".join(f"{arg}\t{datamodule.hparams[arg]}" for arg in datamodule_args))
+    datamodule_args = ["image_width", "image_height", "batch_size", "random_magnitude"]
+    parser.add_argument(
+        "-d", metavar=tuple(map(str.upper, datamodule_args)), nargs=4, dest="datamodule_args", type=int,
+        help="Create new datamodule and exit, current parameters:\n" +
+             "\n".join(f"{arg}\t{datamodule.hparams[arg]}" for arg in datamodule_args))
 
     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", metavar="X", dest="transformer_args", nargs=len(transformer_args),
-                        help="Transformer init args, reference values:\n" +
-                             "\n".join(f"{k}\t{v}" for k, v in transformer_args))
+    parser.add_argument(
+        "-t", metavar=tuple(args[0].upper() for args in transformer_args), dest="transformer_args",
+        nargs=len(transformer_args),
+        help="Transformer init args, default values:\n" + "\n".join(f"{k}\t{v}" for k, v in transformer_args))
 
     args = parser.parse_args()
-
+    if args.data_args:
+        args.data_args = dict(zip(data_args, args.data_args))
     if args.datamodule_args:
         args.datamodule_args = dict(zip(datamodule_args, args.datamodule_args))
 
@@ -63,6 +80,13 @@ def parse_args():
 def main():
     check_setup()
     args = parse_args()
+    if args.data_args:
+        generate_data(examples_count=args.data_args['size'],
+                      max_depth=args.data_args['depth'],
+                      equation_length=args.data_args['length'],
+                      distribution_fraction=args.data_args['fraction'])
+        return
+
     if args.datamodule_args:
         datamodule = LatexImageDataModule(image_width=args.datamodule_args["image_width"],
                                           image_height=args.datamodule_args["image_height"],
@@ -79,23 +103,23 @@ def main():
     logger = None
     callbacks = []
     if args.log:
-        logger = WandbLogger(f"img2tex", log_model=True)
-        callbacks = [LogImageTexCallback(logger, top_k=10, max_length=20),
+        logger = WandbLogger(f"img2tex", save_dir=WANDB_DIR, log_model=True)
+        callbacks = [LogImageTexCallback(logger, top_k=10, max_length=100),
                      LearningRateMonitor(logging_interval="step"),
                      ModelCheckpoint(save_top_k=10,
+                                     every_n_train_steps=500,
                                      monitor="val_loss",
                                      mode="min",
                                      filename="img2tex-{epoch:02d}-{val_loss:.2f}")]
 
-    trainer = Trainer(max_epochs=args.max_epochs,
-                      accelerator="cpu" if args.gpus is None else "gpu",
+    trainer = Trainer(default_root_dir=TRAINER_DIR,
+                      max_epochs=args.max_epochs,
+                      accelerator="gpu" if args.gpus else "cpu",
                       gpus=args.gpus,
                       logger=logger,
                       strategy="ddp_find_unused_parameters_false",
                       enable_progress_bar=True,
-                      default_root_dir=TRAINER_DIR,
-                      callbacks=callbacks,
-                      check_val_every_n_epoch=5)
+                      callbacks=callbacks)
 
     transformer = Transformer(num_encoder_layers=args.transformer_args["num_encoder_layers"],
                               num_decoder_layers=args.transformer_args["num_decoder_layers"],
@@ -111,6 +135,14 @@ def main():
     trainer.fit(transformer, datamodule=datamodule)
     trainer.test(transformer, datamodule=datamodule)
 
+    if args.log:
+        transformer = average_checkpoints(model_type=Transformer, checkpoints_dir=trainer.checkpoint_callback.dirpath)
+        transformer_path = os.path.join(RESOURCES, f"{trainer.logger.version}.pt")
+        transformer.eval()
+        transformer.freeze()
+        torch.save(transformer.state_dict(), transformer_path)
+        print(f"Transformer ensemble saved to '{transformer_path}'")
+
 
 if __name__ == "__main__":
     main()

utils.py

@@ -1,7 +1,7 @@
 from constants import TOKENIZER_PATH
-from data_preprocessing import RandomizeImageTransform
 
 import einops
+import os
 import random
 from pytorch_lightning.callbacks import Callback
 import torch
@@ -22,8 +22,7 @@ class LogImageTexCallback(Callback):
             return
         sample_id = random.randint(0, len(batch['images']) - 1)
         image = batch['images'][sample_id]
-        texs_predicted, texs_ids = beam_search_decode(transformer, image, transform_image=False, top_k=self.top_k,
-                                                      max_length=self.max_length)
+        texs_predicted = beam_search_decode(transformer, image, top_k=self.top_k, max_length=self.max_length)
         image = self.tensor_to_PIL(image)
         tex_true = self.tex_tokenizer.decode(list(batch['tex_ids'][sample_id].to('cpu', torch.int)))
         self.logger.log_image(key="samples", images=[image],
@@ -31,9 +30,8 @@ class LogImageTexCallback(Callback):
 
 
 @torch.inference_mode()
-def beam_search_decode(transformer, image, transform_image=True, top_k=10, max_length=100):
+def beam_search_decode(transformer, image, image_transform=None, top_k=10, max_length=100):
     """Performs decoding maintaining k best candidates"""
-    assert torch.is_tensor(image) and len(image.shape) == 3, "Image must be a 3 dimensional tensor (c h w)"
 
     def get_tgt_padding_mask(tgt):
         mask = tgt == tex_tokenizer.token_to_id("[SEP]")
@@ -41,12 +39,11 @@ def beam_search_decode(transformer, image, transform_image=True, top_k=10, max_l
         mask = mask.to(transformer.device, torch.bool)
         return mask
 
+    if image_transform:
+        image = image_transform(image)
+
+    assert torch.is_tensor(image) and len(image.shape) == 3, "Image must be a 3 dimensional tensor (c h w)"
     src = einops.rearrange(image, "c h w -> () c h w").to(transformer.device)
-    if transform_image:
-        image_transform = RandomizeImageTransform(width=transformer.hparams["image_width"],
-                                                  height=transformer.hparams["image_width"],
-                                                  random_magnitude=0)
-        src = image_transform(src)
     memory = transformer.encode(src)
 
     tex_tokenizer = torch.load(TOKENIZER_PATH)
@@ -82,4 +79,27 @@ def beam_search_decode(transformer, image, transform_image=True, top_k=10, max_l
         padding_mask & (candidates_tex_ids != tex_tokenizer.token_to_id("[SEP]")),
         tex_tokenizer.token_to_id("[PAD]")).tolist()
     texs = tex_tokenizer.decode_batch(candidates_tex_ids, skip_special_tokens=True)
-    return texs, candidates_tex_ids
+    texs = [tex.replace("\\ ", "\\") for tex in texs]
+    return texs
+
+
+def average_checkpoints(model_type, checkpoints_dir):
+    """Returns model averaged from checkpoints
+    Args:
+        :model_type: -- pytorch_lightning.LightningModule that corresponds to checkpoints
+        :checkpoints_dir: -- path to checkpoints
+    """
+    checkpoints = [checkpoint.path for checkpoint in os.scandir(checkpoints_dir)]
+    n_models = len(checkpoints)
+    assert n_models > 0
+    average_model = model_type.load_from_checkpoint(checkpoints[0])
+
+    for checkpoint in checkpoints[1:]:
+        model = model_type.load_from_checkpoint(checkpoint)
+        for weight, weight_to_add in zip(average_model.parameters(), model.parameters()):
+            weight.data.add_(weight_to_add.data)
+
+    for weight in average_model.parameters():
+        weight.data.divide_(n_models)
+
+    return average_model