+
+
+def one_epoch(model, quiz_machine, local_device=main_device):
+ model.to(local_device).train()
+
+ optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)
+
+ nb_train_samples, acc_train_loss = 0, 0.0
+
+ for input in quiz_machine.batches(model, split="train"):
+ input = input.to(local_device)
+
+ if nb_train_samples % args.batch_size == 0:
+ optimizer.zero_grad()
+
+ 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.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} model {model.id} {train_perplexity}")
+
+ run_tests(model, quiz_machine, deterministic_synthesis=False)
+
+ model.to(main_device)
+
+
+######################################################################
+
+
+def standard_validity(logproba):
+ l = logproba.sort(dim=-1).values
+ return (l[:, 0] < math.log(args.proba_not_understands)) & (
+ l[:, 1] > math.log(args.proba_understands)
+ )
+
+
+def valid_quizzes_and_logprobas(recorded, criteria):
+ validated_quizzes, validated_logprobas = [], []
+ for q, lp in recorded:
+ validated_indices = criteria(lp)
+ validated_quizzes.append(q[validated_indices])
+ validated_logprobas.append(lp[validated_indices])
+
+ if len(validated_quizzes) > 0:
+ return torch.cat(validated_quizzes, dim=0), torch.cat(
+ validated_logprobas, dim=0