From bfcef9a8c82ed45528601e85725166241bbee916 Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Sat, 29 Jun 2024 15:14:04 +0300 Subject: [PATCH] Update. --- main.py | 52 ++++++---- mygpt.py | 44 --------- problem.py | 2 +- quizz_machine.py | 253 +++++++++++++++++++++++++++++++++++------------ sky.py | 11 +-- 5 files changed, 228 insertions(+), 134 deletions(-) diff --git a/main.py b/main.py index 232c724..d7fb3d1 100755 --- a/main.py +++ b/main.py @@ -19,7 +19,6 @@ import sky, quizz_machine ###################################################################### -accuracy_to_make_c_quizzes = 0.975 nb_new_c_quizzes_for_train = 1000 nb_new_c_quizzes_for_test = 100 @@ -82,7 +81,15 @@ parser.add_argument("--deterministic_synthesis", action="store_true", default=Fa parser.add_argument("--nb_gpts", type=int, default=5) -parser.add_argument("--nb_correct_to_validate", type=int, default=4) +parser.add_argument("--nb_models_for_generation", type=int, default=1) + +parser.add_argument("--generation_mode", type=str, default="groupthink") + +parser.add_argument("--min_to_validate", type=int, default=4) + +parser.add_argument("--max_to_validate", type=int, default=4) + +parser.add_argument("--accuracy_to_make_c_quizzes", type=float, default=0.975) parser.add_argument("--dirty_debug", action="store_true", default=False) @@ -96,7 +103,7 @@ if args.result_dir is None: ###################################################################### if args.dirty_debug: - accuracy_to_make_c_quizzes = 0.0 + args.accuracy_to_make_c_quizzes = 0.0 nb_new_c_quizzes_for_train = 100 nb_new_c_quizzes_for_test = 10 @@ -371,22 +378,22 @@ def create_c_quizzes( return sum( [ sum([x.size(0) for x in recorded[n]]) - for n in range(args.nb_correct_to_validate, len(models)) + for n in range(args.min_to_validate, args.max_to_validate + 1) ] ) - while nb_validated() < nb_for_train + nb_for_test: - nb_to_validate = nb_for_train + nb_for_test - - if len(model_indexes) == 0: - model_indexes = [i.item() for i in torch.randperm(len(models))] - - model = models[model_indexes.pop()] - - new_c_quizzes, nb_correct, ave_seq_logproba = quizz_machine.create_c_quizzes( - nb=nb_to_validate, - model_for_generation=model, - models_for_validation=models, + nb_to_create = nb_for_train + nb_for_test + + while nb_validated() < nb_to_create: + ( + new_c_quizzes, + nb_correct, + ave_seq_logproba, + ) = quizz_machine.gang_create_c_quizzes( + nb=nb_to_create, + nb_models_for_generation=args.nb_models_for_generation, + models=models, + mode=args.generation_mode, min_ave_seq_logproba=min_ave_seq_logproba, n_epoch=n_epoch, result_dir=args.result_dir, @@ -405,7 +412,7 @@ def create_c_quizzes( recorded[n].append(new_c_quizzes[nb_correct == n].clone()) log_string( - f"keep c_quizzes {nb_validated()*100/nb_generated():.02f}% kept total {nb_validated()} / {nb_to_validate}" + f"keep c_quizzes {nb_validated()*100/nb_generated():.02f}% kept total {nb_validated()} / {nb_to_create}" ) # concatenate and shuffle @@ -418,7 +425,8 @@ def create_c_quizzes( del recorded[n] new_c_quizzes = torch.cat( - [recorded[n] for n in range(args.nb_correct_to_validate, len(models))], dim=0 + [recorded[n] for n in range(args.min_to_validate, args.max_to_validate + 1)], + dim=0, ) new_c_quizzes = new_c_quizzes[ @@ -431,7 +439,11 @@ def create_c_quizzes( quizz_machine.store_c_quizzes(new_c_quizzes[nb_for_train:], for_train=False) for n in recorded.keys(): - s = "_validated" if n >= args.nb_correct_to_validate and n < len(models) else "" + s = ( + "_validated" + if n >= args.min_to_validate and n <= args.max_to_validate + else "" + ) quizz_machine.problem.save_quizzes( recorded[n][:72], args.result_dir, @@ -501,7 +513,7 @@ for n_epoch in range(args.nb_epochs): f"test_set_composition w_quizzes {quizz_machine.nb_batch_w_quizzes} c_quizzes {quizz_machine.nb_batch_c_quizzes}" ) - if min([m.main_test_accuracy for m in models]) >= accuracy_to_make_c_quizzes: + if min([m.main_test_accuracy for m in models]) >= args.accuracy_to_make_c_quizzes: ave_seq_logproba = create_c_quizzes( models, quizz_machine, diff --git a/mygpt.py b/mygpt.py index 7047849..7119c7a 100755 --- a/mygpt.py +++ b/mygpt.py @@ -271,50 +271,6 @@ class MyGPT(nn.Module): bs = self.readout(bs) return bs - # ar_mask is a tensor with 0s and 1s, of same shape as input, with - # 1s where tokens should be generated. The others are kept - # unchanged. - - def masked_inplace_autoregression( - self, - input, - ar_mask, - seq_logproba, - temperature=1.0, - deterministic_synthesis=False, - forbidden_tokens=None, - forced_biases=None, - ): - to_generate = (ar_mask.sum(0) > 0).nonzero() - - if to_generate.min() > 0: - self( - BracketedSequence(input, 0, to_generate.min()) - ) # Needed to initialize the model's cache - for s in range(to_generate.min(), to_generate.max() + 1): - output = self(BracketedSequence(input, s, 1)).x - - logits = output[:, s] - - logits = (logits / temperature).log_softmax(dim=-1) - - if forbidden_tokens is not None: - logits = logits.masked_fill(forbidden_tokens, float("-inf")) - - if forced_biases is not None: - logits = logits + forced_biases[None, :] - - if deterministic_synthesis: - t_next = logits.argmax(-1) - else: - dist = torch.distributions.categorical.Categorical(logits=logits) - t_next = dist.sample() - - all_n = torch.arange(t_next.size(0)) - seq_logproba += logits[all_n, t_next].sum(dim=-1) - - input[:, s] = ar_mask[:, s] * t_next + (1 - ar_mask[:, s]) * input[:, s] - def record_attention(self, v=True): for m in self.modules(): if isinstance(m, QKVAttention): diff --git a/problem.py b/problem.py index 95a9c41..0795de1 100755 --- a/problem.py +++ b/problem.py @@ -9,7 +9,7 @@ class Problem: # returns a nb x (L+1+L) long tensor where L is the length of one # of the two states of a quizz - def generate_seq(self, nb): + def generate_token_sequences(self, nb): pass # save a file to vizualize quizzes, you can save a txt or png file diff --git a/quizz_machine.py b/quizz_machine.py index 8ee0226..239dc68 100755 --- a/quizz_machine.py +++ b/quizz_machine.py @@ -17,6 +17,88 @@ from mygpt import BracketedSequence ###################################################################### +class Gang(nn.Module): + def __init__(self, models, nb_models_for_generation, mode="groupthink"): + super().__init__() + self.models = models + self.nb_models_for_generation = nb_models_for_generation + self.mode = mode + + def forward(self, bs): + # If first = 0, we are re-starting an auto-regressive process, + # that's the right moment to randomize who gonna do it + if bs.first == 0: + self.models_to_use = [ + self.models[k] + for k in torch.randperm(len(self.models))[ + : self.nb_models_for_generation + ] + ] + + all_the_logits = torch.cat( + [model(bs).x[None] for model in self.models_to_use], dim=0 + ) + + if self.mode == "groupthink": + y = all_the_logits.mean(dim=0) + elif self.mode == "groupwork": + m = torch.rand(all_the_logits.size(), device=all_the_logits.device) + m = (m.sort(dim=0).indices == 0).long() + y = (y * m).sum(dim=0) + else: + raise ValueError(f"Invalid mode {self.mode}") + + return BracketedSequence(y, bs.first, bs.nb) + + +###################################################################### + +# ar_mask is a tensor with 0s and 1s, of same shape as input, with +# 1s where tokens should be generated. The others are kept +# unchanged. + + +def one_batch_masked_inplace_autoregression( + model, + input, + ar_mask, + seq_logproba, + temperature=1.0, + deterministic_synthesis=False, + forbidden_tokens=None, + forced_biases=None, +): + to_generate = (ar_mask.sum(0) > 0).nonzero() + + if to_generate.min() > 0: + model( + BracketedSequence(input, 0, to_generate.min()) + ) # Needed to initialize the model's cache + for s in range(to_generate.min(), to_generate.max() + 1): + output = model(BracketedSequence(input, s, 1)).x + + logits = output[:, s] + + logits = (logits / temperature).log_softmax(dim=-1) + + if forbidden_tokens is not None: + logits = logits.masked_fill(forbidden_tokens, float("-inf")) + + if forced_biases is not None: + logits = logits + forced_biases[None, :] + + if deterministic_synthesis: + t_next = logits.argmax(-1) + else: + dist = torch.distributions.categorical.Categorical(logits=logits) + t_next = dist.sample() + + all_n = torch.arange(t_next.size(0)) + seq_logproba += logits[all_n, t_next].sum(dim=-1) + + input[:, s] = ar_mask[:, s] * t_next + (1 - ar_mask[:, s]) * input[:, s] + + def masked_inplace_autoregression( model, batch_size, @@ -51,7 +133,8 @@ def masked_inplace_autoregression( model.eval() for input, ar_mask, seq_logproba in batches: - model.masked_inplace_autoregression( + one_batch_masked_inplace_autoregression( + model=model, input=input, ar_mask=ar_mask, seq_logproba=seq_logproba, @@ -88,8 +171,12 @@ class QuizzMachine: self.batch_size = batch_size self.device = device - self.train_w_quizzes = self.problem.generate_seq(nb_train_samples).to(device) - self.test_w_quizzes = self.problem.generate_seq(nb_test_samples).to(device) + self.train_w_quizzes = self.problem.generate_token_sequences( + nb_train_samples + ).to(device) + self.test_w_quizzes = self.problem.generate_token_sequences(nb_test_samples).to( + device + ) self.nb_codes = max(self.train_w_quizzes.max(), self.test_w_quizzes.max()) + 1 @@ -215,7 +302,7 @@ class QuizzMachine: input = self.train_w_quizzes if for_train else self.test_w_quizzes nb = min(nb, input.size(0)) input[:-nb] = input[nb:].clone() - input[-nb:] = self.problem.generate_seq(nb).to(self.device) + input[-nb:] = self.problem.generate_token_sequences(nb).to(self.device) def store_c_quizzes(self, new_c_quizzes, for_train=True): if for_train: @@ -223,63 +310,7 @@ class QuizzMachine: else: self.test_c_quizzes.append(new_c_quizzes) - def create_c_quizzes( - self, - nb, - model_for_generation, - models_for_validation, - min_ave_seq_logproba, - n_epoch, - result_dir, - logger, - ): - ############################################################### - # Generate quizzes with model - - c_quizzes = torch.empty( - nb, self.train_w_quizzes.size(1), device=self.device, dtype=torch.int64 - ) - - ar_mask = torch.full(c_quizzes.size(), 1, device=self.device) - seq_logproba = torch.empty(ar_mask.size(0), device=self.device) - - temperature = 1 - d_temperature = 1 / 3 - - while True: - seq_logproba[...] = 0 - - masked_inplace_autoregression( - model=model_for_generation, - batch_size=self.batch_size, - input=c_quizzes, - ar_mask=ar_mask, - seq_logproba=seq_logproba, - temperature=temperature, - deterministic_synthesis=False, - # progress_bar_desc="sampling c_quizzes", - device=self.device, - ) - - ave_seq_logproba = seq_logproba.mean() - - if min_ave_seq_logproba is None: - break - - # Oh man that's ugly - if ave_seq_logproba < min_ave_seq_logproba: - if d_temperature > 0: - d_temperature *= -1 / 3 - temperature += d_temperature - elif ave_seq_logproba > min_ave_seq_logproba * 0.99: - if d_temperature < 0: - d_temperature *= -1 / 3 - temperature += d_temperature - else: - break - - logger(f"changing temperature to {temperature}") - + def comput_correctness(self, c_quizzes, models_for_validation): ############################################################### # Create the reverse quizzes @@ -340,6 +371,102 @@ class QuizzMachine: nb_correct.append((correct * reverse_correct)[None, :]) - nb_correct = torch.cat(nb_correct, dim=0).sum(dim=0) + return torch.cat(nb_correct, dim=0).sum(dim=0) + + def generate_quizzes(self, nb, model_for_generation, min_ave_seq_logproba): + ############################################################### + # Generate quizzes with model - return c_quizzes, nb_correct, seq_logproba.mean() + c_quizzes = torch.empty( + nb, self.train_w_quizzes.size(1), device=self.device, dtype=torch.int64 + ) + + ar_mask = torch.full(c_quizzes.size(), 1, device=self.device) + seq_logproba = torch.empty(ar_mask.size(0), device=self.device) + + # bracketing of the temperature to get the target logproba + + temperature = 1 + d_temperature = 1 / 3 + + while True: + seq_logproba[...] = 0 + + masked_inplace_autoregression( + model=model_for_generation, + batch_size=self.batch_size, + input=c_quizzes, + ar_mask=ar_mask, + seq_logproba=seq_logproba, + temperature=temperature, + deterministic_synthesis=False, + # progress_bar_desc="sampling c_quizzes", + device=self.device, + ) + + ave_seq_logproba = seq_logproba.mean() + + # If we do not have target logprobs, get out now + if min_ave_seq_logproba is None: + break + + # Oh man that's ugly + if ave_seq_logproba < min_ave_seq_logproba: + if d_temperature > 0: + d_temperature *= -1 / 3 + temperature += d_temperature + elif ave_seq_logproba > min_ave_seq_logproba * 0.99: + if d_temperature < 0: + d_temperature *= -1 / 3 + temperature += d_temperature + else: + break + + logger(f"changing temperature to {temperature}") + + return c_quizzes, seq_logproba.mean() + + ###################################################################### + + def create_c_quizzes( + self, + nb, + model_for_generation, + models_for_validation, + min_ave_seq_logproba, + n_epoch, + result_dir, + logger, + ): + c_quizzes, ave_seq_logproba = self.generate_quizzes( + nb, model_for_generation, min_ave_seq_logproba + ) + + nb_correct = self.comput_correctness(c_quizzes, models_for_validation) + + return c_quizzes, nb_correct, ave_seq_logproba + + ###################################################################### + + def gang_create_c_quizzes( + self, + nb, + nb_models_for_generation, + models, + mode, + min_ave_seq_logproba, + n_epoch, + result_dir, + logger, + ): + model_for_generation = Gang(models, nb_models_for_generation, mode) + models_for_validation = models + return self.create_c_quizzes( + nb, + model_for_generation, + models_for_validation, + min_ave_seq_logproba, + n_epoch, + result_dir, + logger, + ) diff --git a/sky.py b/sky.py index fdc1689..e93c88a 100755 --- a/sky.py +++ b/sky.py @@ -54,7 +54,7 @@ class Sky(problem.Problem): def direction_tokens(self): return self.token_forward, self.token_backward - def generate_seq(self, nb, return_frame_sequences=False): + def generate_frame_sequences(self, nb): frame_sequences = [] for _ in tqdm.tqdm(range(nb), dynamic_ncols=True, desc="world generation"): @@ -110,11 +110,10 @@ class Sky(problem.Problem): frame_sequences.append(result) - if return_frame_sequences: - return frame_sequences + return frame_sequences - # Randomize the time direction, annd convert to token - # sequences with the time direction tokens added + def generate_token_sequences(self, nb): + frame_sequences = self.generate_frame_sequences(nb) result = [] @@ -260,7 +259,7 @@ if __name__ == "__main__": sky = Sky(height=6, width=8, speed=1, nb_iterations=4) start_time = time.perf_counter() - seq = sky.generate_seq(nb=64) + seq = sky.generate_frame_sequences(nb=64) delay = time.perf_counter() - start_time print(f"{seq.size(0)/delay:02f} seq/s") -- 2.39.5