X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;ds=inline;f=main.py;h=22edf7b79ce01ef71eb4272c46906f6b96c79425;hb=0cd7843a52fd6156b3fc59b1f6c2f86054948ba1;hp=ca0d1524b04d740c8aaa38fa970b503d868bdba8;hpb=d16410119a4e5c1117f7f0fbbe80e3e54f81f28b;p=culture.git diff --git a/main.py b/main.py index ca0d152..22edf7b 100755 --- a/main.py +++ b/main.py @@ -82,12 +82,6 @@ parser.add_argument("--dropout", type=float, default=0.1) parser.add_argument("--deterministic_synthesis", action="store_true", default=False) -parser.add_argument("--no_checkpoint", action="store_true", default=False) - -parser.add_argument("--resume", action="store_true", default=False) - -parser.add_argument("--checkpoint_name", type=str, default="checkpoint.pth") - ############################## # filetask @@ -207,6 +201,12 @@ if args.result_dir is None: ###################################################################### default_task_args = { + "world": { + "model": "37M", + "batch_size": 100, + "nb_train_samples": 250000, + "nb_test_samples": 10000, + }, "file": { "model": "37M", "batch_size": 25, @@ -219,12 +219,6 @@ 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, @@ -677,64 +671,28 @@ log_string(f"vocabulary_size {vocabulary_size}") ############################## -model = mygpt.MyGPT( - vocabulary_size=vocabulary_size, - dim_model=args.dim_model, - dim_keys=args.dim_keys, - dim_hidden=args.dim_hidden, - nb_heads=args.nb_heads, - nb_blocks=args.nb_blocks, - causal=True, - dropout=args.dropout, -) +models = [] + +for k in range(2): + models.append( + mygpt.MyGPT( + vocabulary_size=vocabulary_size, + dim_model=args.dim_model, + dim_keys=args.dim_keys, + dim_hidden=args.dim_hidden, + nb_heads=args.nb_heads, + nb_blocks=args.nb_blocks, + causal=True, + dropout=args.dropout, + ).to(device) + ) -model.to(device) -nb_parameters = sum(p.numel() for p in model.parameters()) +nb_parameters = sum(p.numel() for p in models[0].parameters()) log_string(f"nb_parameters {nb_parameters} ({int(nb_parameters/1e6)}M)") ###################################################################### -nb_epochs_finished = 0 - -if args.no_checkpoint: - log_string(f"not trying to load checkpoint.") - -else: - try: - checkpoint_name = os.path.join(args.result_dir, args.checkpoint_name) - checkpoint = torch.load(checkpoint_name) - nb_epochs_finished = checkpoint["nb_epochs_finished"] - model.load_state_dict(checkpoint["model_state"]) - torch.set_rng_state(checkpoint["rng_state"]) - if torch.cuda.is_available(): - torch.cuda.set_rng_state(checkpoint["cuda_rng_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) - -###################################################################### - -if args.task == "expr" and args.expr_input_file is not None: - task.produce_results( - n_epoch=nb_epochs_finished, - model=model, - result_dir=args.result_dir, - logger=log_string, - deterministic_synthesis=args.deterministic_synthesis, - input_file=args.expr_input_file, - ) - - exit(0) - -###################################################################### - # Compute the entropy of the training tokens token_count = 0 @@ -803,17 +761,6 @@ else: log_string(f"learning_rate_schedule {learning_rate_schedule}") -############################## - -if nb_epochs_finished >= args.nb_epochs: - task.produce_results( - n_epoch=nb_epochs_finished, - model=model, - result_dir=args.result_dir, - logger=log_string, - deterministic_synthesis=args.deterministic_synthesis, - ) - time_pred_result = None ###################################################################### @@ -896,44 +843,64 @@ def run_tests(model, task, deterministic_synthesis): ###################################################################### -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) +def create_quizzes( + other_models, + task, + nb_for_train=1000, + nb_for_test=100, + nb_runs=10, + nb_min_correct=9, + nb_max_correct=9, +): + 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=4 * (nb_for_train + nb_for_test), + models=other_models, + nb_runs=nb_runs, + ) - 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, + to_keep = new_quizzes[ + torch.logical_and( + nb_correct >= nb_min_correct, nb_correct <= nb_max_correct ) + ] + log_string(f"keep {to_keep.size(0)} quizzes") + kept.append(to_keep) - 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] - new_quizzes = torch.cat(kept, dim=0)[: nb_for_train + nb_for_test] + task.store_new_quizzes(new_quizzes[:nb_for_train], for_train=True) + task.store_new_quizzes(new_quizzes[nb_for_train:], for_train=False) - 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, + ) - task.save_image( - new_quizzes[:96], - args.result_dir, - f"world_new_{n_epoch:04d}.png", - log_string, - ) + +###################################################################### + +accuracy_to_make_quizzes = 0.95 + +for n_epoch in range(nb_epochs_finished, args.nb_epochs): + learning_rate = learning_rate_schedule[n_epoch] + + for m in models: + one_epoch(m, task, learning_rate) + test_accuracy = run_tests(m, task, deterministic_synthesis=False) + + if test_accuracy >= accuracy_to_make_quizzes: + other_models = models.copy() + other_models.remove(model) + create_quizzes(other_models, task) # -------------------------------------------- @@ -944,19 +911,4 @@ for n_epoch in range(nb_epochs_finished, args.nb_epochs): ) time_pred_result = time_current_result - # -------------------------------------------- - - checkpoint = { - "nb_epochs_finished": n_epoch + 1, - "model_state": model.state_dict(), - "rng_state": torch.get_rng_state(), - } - - if torch.cuda.is_available(): - checkpoint["cuda_rng_state"] = torch.cuda.get_rng_state() - - checkpoint_name = os.path.join(args.result_dir, args.checkpoint_name) - torch.save(checkpoint, checkpoint_name) - log_string(f"saved checkpoint {checkpoint_name}") - ######################################################################