Spaces:
Runtime error
Runtime error
from constants import TRAINER_DIR, TOKENIZER_PATH, DATAMODULE_PATH | |
from data_preprocessing import LatexImageDataModule | |
from model import Transformer | |
from utils import LogImageTexCallback | |
import argparse | |
import os | |
from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint | |
from pytorch_lightning.loggers import WandbLogger | |
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(): | |
os.environ["TOKENIZERS_PARALLELISM"] = "false" | |
if not os.path.isfile(DATAMODULE_PATH): | |
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): | |
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"] | |
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)) | |
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)) | |
args = parser.parse_args() | |
if args.datamodule_args: | |
args.datamodule_args = dict(zip(datamodule_args, args.datamodule_args)) | |
if args.transformer_args: | |
args.transformer_args = dict(zip(list(zip(*transformer_args))[0], args.transformer_args)) | |
else: | |
args.transformer_args = dict(transformer_args) | |
return args | |
def main(): | |
check_setup() | |
args = parse_args() | |
if args.datamodule_args: | |
datamodule = LatexImageDataModule(image_width=args.datamodule_args["image_width"], | |
image_height=args.datamodule_args["image_height"], | |
batch_size=args.datamodule_args["batch_size"], | |
random_magnitude=args.datamodule_args["random_magnitude"]) | |
datamodule.train_tokenizer() | |
tex_tokenizer = torch.load(TOKENIZER_PATH) | |
print(f"Vocabulary size {tex_tokenizer.get_vocab_size()}") | |
torch.save(datamodule, DATAMODULE_PATH) | |
return | |
datamodule = torch.load(DATAMODULE_PATH) | |
tex_tokenizer = torch.load(TOKENIZER_PATH) | |
logger = None | |
callbacks = [] | |
if args.log: | |
logger = WandbLogger(f"img2tex", log_model=True) | |
callbacks = [LogImageTexCallback(logger, top_k=10, max_length=20), | |
LearningRateMonitor(logging_interval="step"), | |
ModelCheckpoint(save_top_k=10, | |
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", | |
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) | |
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=tex_tokenizer.get_vocab_size(), | |
pad_idx=tex_tokenizer.token_to_id("[PAD]")) | |
trainer.fit(transformer, datamodule=datamodule) | |
trainer.test(transformer, datamodule=datamodule) | |
if __name__ == "__main__": | |
main() | |