From 3b41e2797fc340fd11cb35015b57c3cae1e8447b Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Tue, 2 Jul 2024 12:42:32 +0300 Subject: [PATCH] Update. --- main.py | 6 +- problem.py | 20 +++-- quizz_machine.py | 83 ++++++++++++++----- sky.py | 204 +++++++++++++++++++++-------------------------- 4 files changed, 169 insertions(+), 144 deletions(-) diff --git a/main.py b/main.py index d412e6c..d194a8d 100755 --- a/main.py +++ b/main.py @@ -452,10 +452,8 @@ def create_c_quizzes( q = valid_c_quizzes(recorded, criteria=lambda nb_correct: nb_correct == n)[:72] if q.size(0) > 0: - quizz_machine.problem.save_quizzes( - q, - args.result_dir, - f"culture_c_quiz_{n_epoch:04d}_N{n}{s}", + quizz_machine.save_quizzes( + args.result_dir, f"culture_c_quiz_{n_epoch:04d}_N{n}{s}", q ) diff --git a/problem.py b/problem.py index 0795de1..0bc83a1 100755 --- a/problem.py +++ b/problem.py @@ -7,15 +7,21 @@ 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_token_sequences(self, nb): + def nb_token_values(self): pass - # save a file to vizualize quizzes, you can save a txt or png file - def save_quizzes(self, input, result_dir, filename_prefix): + # returns two tensors nb x D and nb x D' + def generate_prompts_and_answers(self, nb): pass - # returns a pair (forward_tokens, backward_token) - def direction_tokens(self): + # save a file to vizualize quizzes, you can save a txt or png file + def save_quizzes( + self, + result_dir, + filename_prefix, + prompts, + answers, + predicted_prompt=None, + predicted_answers=None, + ): pass diff --git a/quizz_machine.py b/quizz_machine.py index 697f27e..5f19998 100755 --- a/quizz_machine.py +++ b/quizz_machine.py @@ -119,6 +119,20 @@ class QuizzMachine: b = torch.arange(input.size(1), device=input.device) > input.size(1) // 2 return b.long()[None, :].expand_as(input) + def generate_token_sequences(self, nb): + prompts, answers = self.problem.generate_prompts_and_answers(nb) + result = [] + + for prompt, answer in zip(prompts, answers): + if torch.rand(1) < 0.5: + a = [torch.tensor([self.token_forward]), prompt, answer] + else: + a = [torch.tensor([self.token_backward]), answer, prompt] + + result.append(torch.cat(a, dim=0)[None, :]) + + return torch.cat(result, dim=0) + def __init__( self, problem, @@ -131,29 +145,57 @@ class QuizzMachine: ): super().__init__() + v = problem.nb_token_values() + self.token_forward = v + self.token_backward = v + 1 + self.nb_token_values = v + 2 + self.problem = problem self.batch_size = batch_size self.device = device self.logger = logger - 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( + self.train_w_quizzes = self.generate_token_sequences(nb_train_samples).to( device ) - self.nb_codes = max(self.train_w_quizzes.max(), self.test_w_quizzes.max()) + 1 + self.test_w_quizzes = self.generate_token_sequences(nb_test_samples).to(device) self.train_c_quizzes = [] self.test_c_quizzes = [] if result_dir is not None: - self.problem.save_quizzes( - self.train_w_quizzes[:72], result_dir, "culture_w_quizzes" + self.save_quizzes( + result_dir, "culture_w_quizzes", self.train_w_quizzes[:72] ) + def save_quizzes(self, result_dir, filename_prefix, quizzes, prediction=False): + print(f"DEBUG {quizzes.size()=}") + l = (quizzes.size(1) - 1) // 2 + forward = (quizzes[:, 0] == self.token_forward).long() + backward = (quizzes[:, 0] == self.token_backward).long() + assert forward.equal(1 - backward) + first = quizzes[:, 1 : 1 + l] + second = quizzes[:, 1 + l : 1 + 2 * l] + prompts = forward[:, None] * first + backward[:, None] * second + answers = forward[:, None] * second + backward[:, None] * first + + if prediction: + predicted_prompts = backward + predicted_answers = forward + else: + predicted_prompts = None + predicted_answers = None + + self.problem.save_quizzes( + result_dir, + filename_prefix, + prompts, + answers, + predicted_prompts, + predicted_answers, + ) + def batches(self, split="train", desc=None): assert split in {"train", "test"} if split == "train": @@ -194,7 +236,7 @@ class QuizzMachine: yield batch def vocabulary_size(self): - return self.nb_codes + return self.nb_token_values def produce_results( self, n_epoch, model, result_dir, deterministic_synthesis, nmax=1000 @@ -258,8 +300,11 @@ class QuizzMachine: device=self.device, ) - self.problem.save_quizzes( - result[:72], result_dir, f"culture_prediction_{n_epoch:04d}_{model.id:02d}" + self.save_quizzes( + result_dir, + f"culture_prediction_{n_epoch:04d}_{model.id:02d}", + quizzes=result[:72], + prediction=True, ) return main_test_accuracy @@ -268,7 +313,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_token_sequences(nb).to(self.device) + input[-nb:] = self.generate_token_sequences(nb).to(self.device) def store_c_quizzes(self, new_c_quizzes, for_train=True): if for_train: @@ -277,15 +322,15 @@ class QuizzMachine: self.test_c_quizzes.append(new_c_quizzes) def reverse_time(self, c_quizzes): - token_forward, token_backward = self.problem.direction_tokens() - l = (c_quizzes.size(1) - 1) // 2 - direction = c_quizzes[:, l : l + 1] - direction = self.problem.token_forward * ( - direction == self.problem.token_backward - ) + self.problem.token_backward * (direction == self.problem.token_forward) + direction = c_quizzes[:, 0:1] + direction = self.token_forward * ( + direction == self.token_backward + ) + self.token_backward * (direction == self.token_forward) - return torch.cat([c_quizzes[:, l + 1 :], direction, c_quizzes[:, :l]], dim=1) + return torch.cat( + [direction, c_quizzes[:, l + 1 :], c_quizzes[:, 1 : l + 1]], dim=1 + ) def compute_correctness( self, c_quizzes, models_for_validation, both_directions=True diff --git a/sky.py b/sky.py index 4ca4ba7..d2a4568 100755 --- a/sky.py +++ b/sky.py @@ -37,13 +37,14 @@ class Sky(problem.Problem): token_background = 0 first_bird_token = 1 nb_bird_tokens = colors.size(0) - 1 - token_forward = first_bird_token + nb_bird_tokens - token_backward = token_forward + 1 token2char = ( "_" + "".join([chr(ord("A") + n) for n in range(len(colors) - 1)]) + "><" ) + def nb_token_values(self): + return len(self.colors) + def __init__( self, height=6, @@ -60,9 +61,6 @@ class Sky(problem.Problem): self.nb_iterations = nb_iterations self.avoid_collision = avoid_collision - def direction_tokens(self): - return self.token_forward, self.token_backward - def generate_frame_sequences(self, nb): frame_sequences = [] @@ -159,32 +157,11 @@ class Sky(problem.Problem): def generate_prompts_and_answers(self, nb): frame_sequences = self.generate_frame_sequences(nb) - prompts = frame_sequences[:, : frame_sequences.size(0) // 2].flatten(1) - answers = frame_sequences[:, frame_sequences.size(0) // 2 :].flatten(1) + frame_sequences = torch.cat([x[None] for x in frame_sequences], dim=0) + prompts = frame_sequences[:, : frame_sequences.size(1) // 2].flatten(1) + answers = frame_sequences[:, frame_sequences.size(1) // 2 :].flatten(1) return prompts, answers - def generate_token_sequences(self, nb): - frame_sequences = self.generate_frame_sequences(nb) - - result = [] - - for frame_sequence in frame_sequences: - a = [] - if torch.rand(1) < 0.5: - for frame in frame_sequence: - if len(a) > 0: - a.append(torch.tensor([self.token_forward])) - a.append(frame.flatten()) - else: - for frame in reversed(frame_sequence): - if len(a) > 0: - a.append(torch.tensor([self.token_backward])) - a.append(frame.flatten()) - - result.append(torch.cat(a, dim=0)[None, :]) - - return torch.cat(result, dim=0) - ###################################################################### def frame2img(self, x, scale=15): @@ -214,92 +191,82 @@ class Sky(problem.Problem): return x - def seq2img(self, seq, scale=15): - all = [ - self.frame2img( - seq[:, : self.height * self.width].reshape(-1, self.height, self.width), - scale, + def seq2str(self, seq): + result = [] + for s in seq: + result.append("".join([self.token2char[v] for v in s])) + return result + + def save_image( + self, + result_dir, + filename, + prompts, + answers, + predicted_prompts=None, + predicted_answers=None, + ): + if predicted_prompts is None: + predicted_prompts = 255 + + if predicted_answers is None: + predicted_answers = 255 + + def add_frame(x, c, margin): + y = x.new_full( + (x.size(0), x.size(1), x.size(2) + 2 * margin, x.size(3) + 2 * margin), + 0, ) - ] + if type(c) is int: + y[...] = c + else: + c = c.long()[:, None] + c = c * torch.tensor([192, 192, 192], device=c.device) + ( + 1 - c + ) * torch.tensor([255, 255, 255], device=c.device) + y[...] = c[:, :, None, None] + y[:, :, margin:-margin, margin:-margin] = x + return y - separator = torch.full((seq.size(0), 3, self.height * scale - 1, 1), 0) + margin = 4 - t = self.height * self.width + img_prompts = add_frame(self.frame2img(prompts.to("cpu")), 0, 1) + img_answers = add_frame(self.frame2img(answers.to("cpu")), 0, 1) - while t < seq.size(1): - direction_tokens = seq[:, t] - t += 1 + # img_prompts = add_frame(img_prompts, 255, margin) + # img_answers = add_frame(img_answers, 255, margin) - direction_images = self.colors[ - torch.full( - (direction_tokens.size(0), self.height * scale - 1, scale), 0 - ) - ].permute(0, 3, 1, 2) - - for n in range(direction_tokens.size(0)): - if direction_tokens[n] == self.token_forward: - for k in range(scale): - for l in [0, 1]: - direction_images[ - n, - :, - (self.height * scale) // 2 - scale // 2 + k - l, - 3 + scale // 2 - abs(k - scale // 2), - ] = 0 - elif direction_tokens[n] == self.token_backward: - for k in range(scale): - for l in [0, 1]: - direction_images[ - n, - :, - (self.height * scale) // 2 - scale // 2 + k - l, - 3 + abs(k - scale // 2), - ] = 0 - else: - for k in range(2, scale - 2): - for l in [0, 1]: - direction_images[ - n, - :, - (self.height * scale) // 2 - scale // 2 + k - l, - k, - ] = 0 - direction_images[ - n, - :, - (self.height * scale) // 2 - scale // 2 + k - l, - scale - 1 - k, - ] = 0 - - all += [ - separator, - direction_images, - separator, - self.frame2img( - seq[:, t : t + self.height * self.width].reshape( - -1, self.height, self.width - ), - scale, - ), - ] - - t += self.height * self.width - - return torch.cat(all, dim=3) + img_prompts = add_frame(img_prompts, predicted_prompts, margin) + img_answers = add_frame(img_answers, predicted_answers, margin) - def seq2str(self, seq): - result = [] - for s in seq: - result.append("".join([self.token2char[v] for v in s])) - return result + separator = img_prompts.new_full( + (img_prompts.size(0), img_prompts.size(1), img_prompts.size(2), margin), 255 + ) + + img = torch.cat([img_prompts, img_answers], dim=3) - def save_image(self, input, result_dir, filename): - img = self.seq2img(input.to("cpu")) image_name = os.path.join(result_dir, filename) - torchvision.utils.save_image(img.float() / 255.0, image_name, nrow=6, padding=4) + torchvision.utils.save_image( + img.float() / 255.0, image_name, nrow=6, padding=margin * 2, pad_value=1.0 + ) - def save_quizzes(self, input, result_dir, filename_prefix): - self.save_image(input, result_dir, filename_prefix + ".png") + def save_quizzes( + self, + result_dir, + filename_prefix, + prompts, + answers, + predicted_prompts=None, + predicted_answers=None, + ): + self.save_image( + result_dir, + filename_prefix + ".png", + prompts, + answers, + predicted_prompts, + predicted_answers, + ) ###################################################################### @@ -309,10 +276,19 @@ if __name__ == "__main__": sky = Sky(height=6, width=8, speed=4, nb_iterations=2) - start_time = time.perf_counter() - token_sequences = sky.generate_token_sequences(nb=64) - delay = time.perf_counter() - start_time - print(f"{token_sequences.size(0)/delay:02f} seq/s") + prompts, answers = sky.generate_prompts_and_answers(4) + + predicted_prompts = torch.rand(prompts.size(0)) < 0.5 + predicted_answers = torch.rand(answers.size(0)) < 0.5 + + sky.save_quizzes( + "/tmp", "test", prompts, answers, predicted_prompts, predicted_answers + ) + + # start_time = time.perf_counter() + # token_sequences = sky.generate_token_sequences(nb=64) + # delay = time.perf_counter() - start_time + # print(f"{token_sequences.size(0)/delay:02f} seq/s") # print(sky.seq2str(seq[:4])) @@ -330,9 +306,9 @@ if __name__ == "__main__": # seq = (1 - m) * seq + m * 23 # print(seq.size()) - img = sky.seq2img(token_sequences) + # img = sky.seq2img(token_sequences) # print(img.size()) - torchvision.utils.save_image( - img.float() / 255.0, "/tmp/world.png", nrow=6, padding=6, pad_value=0 - ) + # torchvision.utils.save_image( + # img.float() / 255.0, "/tmp/world.png", nrow=6, padding=6, pad_value=0 + # ) -- 2.39.5