fixed image transform

data_preprocessing.py

@@ -74,7 +74,7 @@ class RandomizeImageTransform(object):
 
     def __init__(self, width, height, random_magnitude):
         self.transform = T.Compose((
-            lambda x: x if random_magnitude == 0 else T.ColorJitter(brightness=random_magnitude / 10,
+            (lambda x: x) if random_magnitude == 0 else T.ColorJitter(brightness=random_magnitude / 10,
                                                                     contrast=random_magnitude / 10,
                                                                     saturation=random_magnitude / 10,
                                                                     hue=min(0.5, random_magnitude / 10)),
@@ -83,7 +83,7 @@ class RandomizeImageTransform(object):
             T.functional.invert,
             T.CenterCrop((height, width)),
             torch.Tensor.contiguous,
-            lambda x: x if random_magnitude == 0 else T.RandAugment(magnitude=random_magnitude),
+            (lambda x: x) if random_magnitude == 0 else T.RandAugment(magnitude=random_magnitude),
             T.ConvertImageDtype(torch.float32)
         ))
 

train.py

@@ -13,8 +13,7 @@ import torch
 
 
 def check_setup():
-    print(
-        "Disabling tokenizers parallelism because it can't be used before forking and I didn't bother to figure it out")
+    # 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")
@@ -107,7 +106,7 @@ def main():
         callbacks = [LogImageTexCallback(logger, top_k=10, max_length=100),
                      LearningRateMonitor(logging_interval="step"),
                      ModelCheckpoint(save_top_k=10,
-                                     every_n_train_steps=500,
+                                     every_n_train_steps=5,
                                      monitor="val_loss",
                                      mode="min",
                                      filename="img2tex-{epoch:02d}-{val_loss:.2f}")]
@@ -135,9 +134,9 @@ def main():
     trainer.fit(transformer, datamodule=datamodule)
     trainer.test(transformer, datamodule=datamodule)
 
-    if args.log:
+    if args.log and len(os.listdir(trainer.checkpoint_callback.dirpath)):
         transformer = average_checkpoints(model_type=Transformer, checkpoints_dir=trainer.checkpoint_callback.dirpath)
-        transformer_path = os.path.join(RESOURCES, f"{trainer.logger.version}.pt")
+        transformer_path = os.path.join(RESOURCES, f"model_{trainer.logger.version}.pt")
         transformer.eval()
         transformer.freeze()
         torch.save(transformer.state_dict(), transformer_path)