From cd3329fc206bacfd90a8e2cbe364244359568733 Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Sun, 10 Dec 2023 16:00:57 +0100 Subject: [PATCH] Update. --- main.py | 25 ++++++++++++++++- problems.py | 81 +++++++++++++++++++++++++++++++++++++++++++++-------- tasks.py | 50 ++++++++++++++++++--------------- 3 files changed, 122 insertions(+), 34 deletions(-) diff --git a/main.py b/main.py index 17936c3..1d52b6d 100755 --- a/main.py +++ b/main.py @@ -33,7 +33,7 @@ parser.add_argument( "--task", type=str, default="twotargets", - help="byheart, learnop, guessop, mixing, twotargets, addition, picoclvr, mnist, maze, snake, stack, expr, rpl, grid, qmlp", + help="byheart, learnop, guessop, mixing, memory, twotargets, addition, picoclvr, mnist, maze, snake, stack, expr, rpl, grid, qmlp", ) parser.add_argument("--log_filename", type=str, default="train.log", help=" ") @@ -256,6 +256,12 @@ default_task_args = { "nb_train_samples": 50000, "nb_test_samples": 10000, }, + "memory": { + "model": "4M", + "batch_size": 100, + "nb_train_samples": 5000, + "nb_test_samples": 1000, + }, "mixing": { "model": "37M", "batch_size": 25, @@ -285,6 +291,13 @@ default_model_args = { "nb_heads": 2, "nb_blocks": 2, }, + "4M": { + "dim_model": 256, + "dim_keys": 32, + "dim_hidden": 1024, + "nb_heads": 4, + "nb_blocks": 6, + }, "37M": { "dim_model": 512, "dim_keys": 64, @@ -418,6 +431,16 @@ elif args.task == "twotargets": device=device, ) +elif args.task == "memory": + task = tasks.SandBox( + problem=problems.ProblemMemory(), + nb_train_samples=args.nb_train_samples, + nb_test_samples=args.nb_test_samples, + batch_size=args.batch_size, + logger=log_string, + device=device, + ) + elif args.task == "mixing": task = tasks.SandBox( problem=problems.ProblemMixing( diff --git a/problems.py b/problems.py index ac16df4..d7dbc54 100755 --- a/problems.py +++ b/problems.py @@ -110,6 +110,48 @@ class ProblemDegradation(Problem): #################### +class ProblemMemory(Problem): + def __init__(self, len_total=25): + self.len_total = len_total + self.max_len_pattern = 5 + self.nb_noise_tokens = 10 + self.start_pattern_token = 0 + self.end_pattern_token = 1 + self.start_result_token = 2 + self.end_result_token = 3 + self.token_string = "[]<>" + "".join( + [chr(ord("a") + k) for k in range(self.nb_noise_tokens)] + ) + + def generate_sequences(self, nb): + sequences = ( + torch.randint(self.nb_noise_tokens, (nb, self.len_total)) + + self.end_result_token + + 1 + ) + len_patterns = torch.randint(self.max_len_pattern, (nb,)) + 1 + pattern_positions = torch.randint( + self.len_total - (5 + 2 * self.max_len_pattern), (nb,) + ) + k = self.len_total - (3 + self.max_len_pattern) + for i in range(nb): + l = len_patterns[i] + j = pattern_positions[i] + sequences[i, j] = self.start_pattern_token + sequences[i, j + l + 2] = self.end_pattern_token + sequences[i, k] = self.start_result_token + sequences[i, k + l + 2] = self.end_result_token + sequences[i, k + 1 : k + 2 + l] = sequences[i, j + 1 : j + 2 + l] + + j = torch.arange(self.len_total)[None, :] + ar_mask = (j > k).long() * (j <= k + 1 + len_patterns[:, None]).long() + + return sequences, ar_mask + + def seq2str(self, seq): + return "".join(self.token_string[x.item()] for x in seq) + + class ProblemTwoTargets(Problem): def __init__(self, len_total=10, len_targets=3): assert len_targets >= 3 @@ -325,22 +367,38 @@ class ProblemMixing(Problem): return y def start_error(self, x): - i = torch.arange(self.height, device=x.device).reshape(1, -1, 1).expand_as(x) - j = torch.arange(self.width, device=x.device).reshape(1, 1, -1).expand_as(x) - - ri = ( - (x == self.height * self.width).long().sum(dim=-1).argmax(-1).view(-1, 1, 1) - ) - rj = ( - (x == self.height * self.width).long().sum(dim=-2).argmax(-1).view(-1, 1, 1) - ) + if self.random_start: + i = ( + torch.arange(self.height, device=x.device) + .reshape(1, -1, 1) + .expand_as(x) + ) + j = torch.arange(self.width, device=x.device).reshape(1, 1, -1).expand_as(x) + + ri = ( + (x == self.height * self.width) + .long() + .sum(dim=-1) + .argmax(-1) + .view(-1, 1, 1) + ) + rj = ( + (x == self.height * self.width) + .long() + .sum(dim=-2) + .argmax(-1) + .view(-1, 1, 1) + ) - m = 1 - torch.logical_or(i == ri, j == rj).long().flatten(1) + m = 1 - torch.logical_or(i == ri, j == rj).long().flatten(1) + else: + m = 1 x = x.flatten(1) u = torch.arange(self.height * self.width, device=x.device).reshape(1, -1) d = (x - (m * u + (1 - m) * self.height * self.width)).abs().sum(-1) + return d def moves(self, x): @@ -424,7 +482,8 @@ class ProblemMixing(Problem): #################### if __name__ == "__main__": - p = ProblemMixing() + p = ProblemMixing(height=3, width=3, random_start=False) + s, m = p.generate_sequences(10000) for x in s[:5]: print(p.seq2str(x)) diff --git a/tasks.py b/tasks.py index 7a4abbe..f4be293 100755 --- a/tasks.py +++ b/tasks.py @@ -125,6 +125,12 @@ class SandBox(Task): (0, 1), } + if logger is not None: + for s, a in zip(self.train_input[:100], self.train_ar_mask[:100]): + logger(f"train_sequences {self.problem.seq2str(s)}") + a = "".join(["01"[x.item()] for x in a]) + logger(f" {a}") + def batches(self, split="train", nb_to_use=-1, desc=None): assert split in {"train", "test"} input = self.train_input if split == "train" else self.test_input @@ -206,30 +212,30 @@ class SandBox(Task): with torch.autograd.no_grad(): t = model.training model.eval() - model.record_attention(True) + # model.record_attention(True) model(BracketedSequence(input)) model.train(t) - ram = model.retrieve_attention() - model.record_attention(False) - - tokens_output = [c for c in self.problem.seq2str(input[0])] - tokens_input = ["n/a"] + tokens_output[:-1] - for n_head in range(ram[0].size(1)): - filename = os.path.join( - result_dir, f"sandbox_attention_{k}_h{n_head}.pdf" - ) - attention_matrices = [m[0, n_head] for m in ram] - save_attention_image( - filename, - tokens_input, - tokens_output, - attention_matrices, - k_top=10, - # min_total_attention=0.9, - token_gap=12, - layer_gap=50, - ) - logger(f"wrote {filename}") + # ram = model.retrieve_attention() + # model.record_attention(False) + + # tokens_output = [c for c in self.problem.seq2str(input[0])] + # tokens_input = ["n/a"] + tokens_output[:-1] + # for n_head in range(ram[0].size(1)): + # filename = os.path.join( + # result_dir, f"sandbox_attention_{k}_h{n_head}.pdf" + # ) + # attention_matrices = [m[0, n_head] for m in ram] + # save_attention_image( + # filename, + # tokens_input, + # tokens_output, + # attention_matrices, + # k_top=10, + ##min_total_attention=0.9, + # token_gap=12, + # layer_gap=50, + # ) + # logger(f"wrote {filename}") ###################################################################### -- 2.39.5