Update.

[mygpt.git] / main.py

diff --git a/main.py b/main.py
index 77c4b9e..6c1def7 100755 (executable)
--- a/main.py
+++ b/main.py
@@ -69,6 +69,9 @@ parser.add_argument('--dropout',
 parser.add_argument('--synthesis_sampling',
                     action='store_true', default = True)
 
+parser.add_argument('--no_checkpoint',
+                    action='store_true', default = False)
+
 parser.add_argument('--checkpoint_name',
                     type = str, default = 'checkpoint.pth')
 
@@ -152,6 +155,7 @@ class TaskPicoCLVR(Task):
         self.train_descr = generate_descr((nb * 4) // 5)
         self.test_descr = generate_descr((nb * 1) // 5)
 
+        # Build the tokenizer
         tokens = set()
         for d in [ self.train_descr, self.test_descr ]:
             for s in d:
@@ -212,12 +216,12 @@ class TaskPicoCLVR(Task):
 
         img = [ picoclvr.descr2img(d, height = self.height, width = self.width) for d in descr ]
         img = torch.cat(img, 0)
-        file_name = f'result_picoclvr_{n_epoch:04d}.png'
+        image_name = f'result_picoclvr_{n_epoch:04d}.png'
         torchvision.utils.save_image(
             img / 255.,
-            file_name, nrow = nb_per_primer, pad_value = 0.8
+            image_name, nrow = nb_per_primer, pad_value = 0.8
         )
-        log_string(f'wrote {file_name}')
+        log_string(f'wrote {image_name}')
 
         nb_missing = sum( [
             x[2] for x in picoclvr.nb_missing_properties(
@@ -429,19 +433,23 @@ else:
 
 nb_epochs_finished = 0
 
-try:
-    checkpoint = torch.load(args.checkpoint_name, map_location = device)
-    nb_epochs_finished = checkpoint['nb_epochs_finished']
-    model.load_state_dict(checkpoint['model_state'])
-    optimizer.load_state_dict(checkpoint['optimizer_state'])
-    print(f'Checkpoint loaded with {nb_epochs_finished} epochs finished.')
-
-except FileNotFoundError:
-    print('Starting from scratch.')
+if args.no_checkpoint:
+    log_string(f'Not trying to load checkpoint.')
 
-except:
-    print('Error when loading the checkpoint.')
-    exit(1)
+else:
+    try:
+        checkpoint = torch.load(args.checkpoint_name, map_location = device)
+        nb_epochs_finished = checkpoint['nb_epochs_finished']
+        model.load_state_dict(checkpoint['model_state'])
+        optimizer.load_state_dict(checkpoint['optimizer_state'])
+        log_string(f'Checkpoint loaded with {nb_epochs_finished} epochs finished.')
+
+    except FileNotFoundError:
+        log_string('Starting from scratch.')
+
+    except:
+        log_string('Error when loading the checkpoint.')
+        exit(1)
 
 ######################################################################