Update.

[culture.git] / main.py

diff --git a/main.py b/main.py
index 87a67c3..8715711 100755 (executable)
--- a/main.py
+++ b/main.py
@@ -443,11 +443,15 @@ def create_c_quizzes(
     q = new_c_quizzes[:72]
 
     if q.size(0) > 0:
-        quiz_machine.save_quizzes(args.result_dir, f"culture_c_quiz_{n_epoch:04d}", q)
+        quiz_machine.save_quiz_illustrations(
+            args.result_dir, f"culture_c_quiz_{n_epoch:04d}", q
+        )
 
 
 ######################################################################
 
+nb_loaded_models = 0
+
 models = []
 
 for k in range(args.nb_gpts):
@@ -471,8 +475,23 @@ for k in range(args.nb_gpts):
     model.test_w_quizzes = quiz_machine.generate_token_sequences(args.nb_test_samples)
     quiz_machine.reverse_random_half_in_place(model.test_w_quizzes)
 
+    filename = f"gpt_{model.id:03d}.pth"
+
+    try:
+        model.load_state_dict(torch.load(os.path.join(args.result_dir, filename)))
+        log_string(f"model {model.id} successfully loaded from checkpoint.")
+        nb_loaded_models += 1
+
+    except FileNotFoundError:
+        log_string(f"starting model {model.id} from scratch.")
+
+    except:
+        log_string(f"error when loading {filename}.")
+        exit(1)
+
     models.append(model)
 
+assert nb_loaded_models == 0 or nb_loaded_models == len(models)
 
 nb_parameters = sum(p.numel() for p in models[0].parameters())
 log_string(f"nb_parameters {nb_parameters} ({int(nb_parameters/1e6)}M)")
@@ -547,27 +566,6 @@ if args.dirty_debug:
         return l[:, 0] < math.log(0.5)
 
 
-######################################################################
-
-nb_loaded_models = 0
-
-for model in models:
-    filename = f"gpt_{model.id:03d}.pth"
-
-    try:
-        model.load_state_dict(torch.load(os.path.join(args.result_dir, filename)))
-        log_string(f"model {model.id} successfully loaded from checkpoint.")
-        nb_loaded_models += 1
-
-    except FileNotFoundError:
-        log_string(f"starting model {model.id} from scratch.")
-
-    except:
-        log_string(f"error when loading {filename}.")
-        exit(1)
-
-assert nb_loaded_models == 0 or nb_loaded_models == len(models)
-
 ######################################################################
 
 for n_epoch in range(args.nb_epochs):