Update.
[culture.git] / main.py
diff --git a/main.py b/main.py
index dace5f2..ca0d152 100755 (executable)
--- a/main.py
+++ b/main.py
@@ -219,6 +219,12 @@ default_task_args = {
         "nb_train_samples": 250000,
         "nb_test_samples": 10000,
     },
+    "world": {
+        "model": "37M",
+        "batch_size": 25,
+        "nb_train_samples": 50000,
+        "nb_test_samples": 10000,
+    },
     "byheart": {
         "model": "37M",
         "batch_size": 25,
@@ -463,6 +469,17 @@ elif args.task == "byheart":
     )
     args.max_percents_of_test_in_train = -1
 
+elif args.task == "world":
+    task = tasks.World(
+        nb_train_samples=args.nb_train_samples,
+        nb_test_samples=args.nb_test_samples,
+        batch_size=args.physical_batch_size,
+        result_dir=args.result_dir,
+        logger=log_string,
+        device=device,
+    )
+    args.max_percents_of_test_in_train = -1
+
 elif args.task == "learnop":
     task = tasks.SandBox(
         problem=problems.ProblemLearnOperator(),
@@ -799,9 +816,10 @@ if nb_epochs_finished >= args.nb_epochs:
 
 time_pred_result = None
 
-for n_epoch in range(nb_epochs_finished, args.nb_epochs):
-    learning_rate = learning_rate_schedule[n_epoch]
+######################################################################
+
 
+def one_epoch(model, task, learning_rate):
     log_string(f"learning_rate {learning_rate}")
 
     if args.optim == "sgd":
@@ -815,7 +833,7 @@ for n_epoch in range(nb_epochs_finished, args.nb_epochs):
 
     model.train()
 
-    nb_train_samples, acc_train_loss_ar, acc_train_loss_ae = 0, 0.0, 0.0
+    nb_train_samples, acc_train_loss = 0, 0.0
 
     for input in task.batches(split="train"):
         input = input.to(device)
@@ -823,83 +841,110 @@ for n_epoch in range(nb_epochs_finished, args.nb_epochs):
         if nb_train_samples % args.batch_size == 0:
             optimizer.zero_grad()
 
-        if args.autoencoder_weight > 0:
-            bs_ar, bs_ae = model(mygpt.BracketedSequence(input), autoencoder=True)
-            output_ar, output_ae = bs_ar.x, bs_ae.x
-            loss_ar = F.cross_entropy(output_ar.transpose(1, 2), input)
-            loss_ae = F.cross_entropy(output_ae[:, 1:].transpose(1, 2), input[:, :-1])
-        else:
-            output = model(mygpt.BracketedSequence(input)).x
-            loss_ar = F.cross_entropy(output.transpose(1, 2), input)
-            loss_ae = loss_ar.new_full((1,), 0.0)
-
-        acc_train_loss_ar += loss_ar.item() * input.size(0)
-        acc_train_loss_ae += loss_ae.item() * input.size(0)
+        output = model(mygpt.BracketedSequence(input)).x
+        loss = F.cross_entropy(output.transpose(1, 2), input)
+        acc_train_loss += loss.item() * input.size(0)
 
         nb_train_samples += input.size(0)
 
-        (loss_ar + args.autoencoder_weight * loss_ae).backward()
+        loss.backward()
 
         if nb_train_samples % args.batch_size == 0:
             optimizer.step()
 
+    train_perplexity = math.exp(min(100, acc_train_loss / nb_train_samples))
+
+    log_string(f"train_perplexity {n_epoch} {train_perplexity}")
+
+
+######################################################################
+
+
+def run_tests(model, task, deterministic_synthesis):
     with torch.autograd.no_grad():
         model.eval()
 
-        nb_test_samples, acc_test_loss_ar, acc_test_loss_ae = 0, 0.0, 0.0
+        nb_test_samples, acc_test_loss = 0, 0.0
         nb_samples_accumulated = 0
 
         for input in task.batches(split="test"):
             input = input.to(device)
 
-            if args.autoencoder_weight > 0:
-                bs_ar, bs_ae = model(mygpt.BracketedSequence(input), autoencoder=True)
-                output_ar, output_ae = bs_ar.x, bs_ae.x
-                loss_ae = F.cross_entropy(
-                    output_ae[:, 1:].transpose(1, 2), input[:, :-1]
-                )
-                acc_test_loss_ae += loss_ae.item() * input.size(0)
-            else:
-                bs_ar = model(mygpt.BracketedSequence(input))
-                output_ar = bs_ar.x
+            bs = model(mygpt.BracketedSequence(input))
+            output = bs.x
 
-            loss_ar = F.cross_entropy(output_ar.transpose(1, 2), input)
+            loss = F.cross_entropy(output.transpose(1, 2), input)
 
-            acc_test_loss_ar += loss_ar.item() * input.size(0)
+            acc_test_loss += loss.item() * input.size(0)
 
             nb_test_samples += input.size(0)
 
-        train_ar_perplexity = math.exp(min(100, acc_train_loss_ar / nb_train_samples))
-        test_ar_perplexity = math.exp(min(100, acc_test_loss_ar / nb_test_samples))
-
-        log_string(
-            f"perplexity_ar {n_epoch} train_set {train_set_perplexity} train_prediction {train_ar_perplexity} test_prediction {test_ar_perplexity}"
-        )
-
-        if args.autoencoder_weight > 0:
-            train_ae_perplexity = math.exp(
-                min(100, acc_train_loss_ae / nb_train_samples)
-            )
-            test_ae_perplexity = math.exp(min(100, acc_test_loss_ae / nb_test_samples))
-
-            log_string(
-                f"perplexity_ae {n_epoch} train_set {train_set_perplexity} train_prediction {train_ae_perplexity} test_prediction {test_ae_perplexity}"
-            )
-
-        task.produce_results(
+        main_test_accuracy = task.produce_results(
             n_epoch=n_epoch,
             model=model,
             result_dir=args.result_dir,
             logger=log_string,
-            deterministic_synthesis=args.deterministic_synthesis,
+            deterministic_synthesis=deterministic_synthesis,
         )
 
-        time_current_result = datetime.datetime.now()
-        if time_pred_result is not None:
-            log_string(
-                f"next_result {time_current_result + (time_current_result - time_pred_result)}"
+        test_perplexity = math.exp(min(100, acc_test_loss / nb_test_samples))
+
+        log_string(f"test_perplexity {n_epoch} {test_perplexity}")
+
+    return main_test_accuracy
+
+
+######################################################################
+
+for n_epoch in range(nb_epochs_finished, args.nb_epochs):
+    learning_rate = learning_rate_schedule[n_epoch]
+
+    one_epoch(model, task, learning_rate)
+
+    test_accuracy = run_tests(model, task, deterministic_synthesis=False)
+
+    # --------------------------------------------
+
+    if test_accuracy >= 0.8:
+        nb_for_train, nb_for_test = 1000, 100
+        kept = []
+
+        while sum([x.size(0) for x in kept]) < nb_for_train + nb_for_test:
+            new_quizzes, nb_correct = task.create_new_quizzes(
+                n_epoch=n_epoch,
+                result_dir=args.result_dir,
+                logger=log_string,
+                nb=nb_required,
+                model=model,
+                nb_runs=10,
             )
-        time_pred_result = time_current_result
+
+            to_keep = new_quizzes[torch.logical_and(nb_correct >= 8, nb_correct < 10)]
+            log_string(f"keep {to_keep.size(0)} quizzes")
+            kept.append(to_keep)
+
+        new_quizzes = torch.cat(kept, dim=0)[: nb_for_train + nb_for_test]
+
+        task.store_new_quizzes(new_quizzes[:nb_for_train], train=True)
+        task.store_new_quizzes(new_quizzes[nb_for_train:], train=False)
+
+        task.save_image(
+            new_quizzes[:96],
+            args.result_dir,
+            f"world_new_{n_epoch:04d}.png",
+            log_string,
+        )
+
+    # --------------------------------------------
+
+    time_current_result = datetime.datetime.now()
+    if time_pred_result is not None:
+        log_string(
+            f"next_result {time_current_result + (time_current_result - time_pred_result)}"
+        )
+    time_pred_result = time_current_result
+
+    # --------------------------------------------
 
     checkpoint = {
         "nb_epochs_finished": n_epoch + 1,