From 90eab15841632ef4f7bd88d2a7cbbb2426bf736a Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Sun, 23 Jun 2024 08:43:45 +0200 Subject: [PATCH] Update. --- main.py | 2 +- tasks.py | 4 +- world.py | 173 +++++++++++++++++++++++++++---------------------------- 3 files changed, 89 insertions(+), 90 deletions(-) diff --git a/main.py b/main.py index 79b4b56..3b29d01 100755 --- a/main.py +++ b/main.py @@ -404,7 +404,7 @@ if args.check: nb_new_quizzes_for_test = 10 for n_epoch in range(args.nb_epochs): - a = [(model.id, model.main_test_accuracy) for model in models] + a = [(model.id, float(model.main_test_accuracy)) for model in models] a.sort(key=lambda p: p[0]) log_string(f"current accuracies {a}") diff --git a/tasks.py b/tasks.py index ad95237..1254323 100755 --- a/tasks.py +++ b/tasks.py @@ -104,11 +104,11 @@ class World(Task): self.height = 7 self.width = 9 - self.train_input = world.generate( + self.train_input = world.generate_seq( nb_train_samples, height=self.height, width=self.width ).to(device) - self.test_input = world.generate( + self.test_input = world.generate_seq( nb_test_samples, height=self.height, width=self.width ).to(device) diff --git a/world.py b/world.py index 68f46de..0d3509f 100755 --- a/world.py +++ b/world.py @@ -18,26 +18,16 @@ from torch.nn import functional as F colors = torch.tensor( [ [255, 255, 255], - [255, 20, 147], - [0, 0, 255], + [255, 0, 0], [0, 192, 0], + [0, 0, 255], + [255, 192, 0], [0, 255, 255], + [255, 0, 255], + [192, 255, 192], + [255, 192, 192], + [192, 192, 255], [192, 192, 192], - [106, 90, 205], - [255, 0, 0], - [220, 20, 60], - [65, 105, 225], - [255, 200, 0], - # [255, 182, 193], - # [75, 0, 130], - # [128, 0, 128], - # [30, 144, 255], - # [135, 206, 235], - # [0, 255, 0], - # [64, 224, 208], - # [250, 128, 114], - # [255, 165, 0], - # [0, 255, 255], ] ) @@ -50,87 +40,93 @@ token_backward = token_forward + 1 token2char = "_" + "".join([chr(ord("A") + n) for n in range(len(colors) - 1)]) + "><" -def generate( +def generate_seq( nb, height, width, nb_birds=3, - nb_iterations=1, + nb_iterations=2, ): pairs = [] for _ in tqdm.tqdm(range(nb), dynamic_ncols=True, desc="world generation"): - f_start = torch.zeros(height, width, dtype=torch.int64) + while True: + f_start = torch.zeros(height, width, dtype=torch.int64) + + i, j, vi, vj = ( + torch.empty(nb_birds, dtype=torch.int64), + torch.empty(nb_birds, dtype=torch.int64), + torch.empty(nb_birds, dtype=torch.int64), + torch.empty(nb_birds, dtype=torch.int64), + ) + + col = torch.randperm(colors.size(0) - 1)[:nb_birds].sort().values + 1 - i, j, vi, vj = ( - torch.empty(nb_birds, dtype=torch.int64), - torch.empty(nb_birds, dtype=torch.int64), - torch.empty(nb_birds, dtype=torch.int64), - torch.empty(nb_birds, dtype=torch.int64), - ) - - col = torch.randperm(colors.size(0) - 1)[:nb_birds].sort().values + 1 - - for n in range(nb_birds): - c = col[n] - - while True: - i[n], j[n] = ( - torch.randint(height, (1,))[0], - torch.randint(width, (1,))[0], - ) - vm = torch.randint(4, (1,))[0] - vi[n], vj[n] = (vm % 2) * 2 - 1, (vm // 2) * 2 - 1 - if ( - i[n] - vi[n] >= 0 - and i[n] - vi[n] < height - and j[n] - vj[n] >= 0 - and j[n] - vj[n] < width - and f_start[i[n], j[n]] == 0 - and f_start[i[n] - vi[n], j[n]] == 0 - and f_start[i[n], j[n] - vj[n]] == 0 - ): - break - - f_start[i[n], j[n]] = c - f_start[i[n] - vi[n], j[n]] = c - f_start[i[n], j[n] - vj[n]] = c - - f_end = f_start.clone() - - for l in range(nb_iterations): for n in range(nb_birds): c = col[n] - f_end[i[n], j[n]] = 0 - f_end[i[n] - vi[n], j[n]] = 0 - f_end[i[n], j[n] - vj[n]] = 0 - pi, pj, pvi, pvj = i[n].item(), j[n].item(), vi[n].item(), vj[n].item() + while True: + i[n], j[n] = ( + torch.randint(height, (1,))[0], + torch.randint(width, (1,))[0], + ) + vm = torch.randint(4, (1,))[0] + vi[n], vj[n] = (vm % 2) * 2 - 1, (vm // 2) * 2 - 1 + if ( + i[n] - vi[n] >= 0 + and i[n] - vi[n] < height + and j[n] - vj[n] >= 0 + and j[n] - vj[n] < width + and f_start[i[n], j[n]] == 0 + and f_start[i[n] - vi[n], j[n]] == 0 + and f_start[i[n], j[n] - vj[n]] == 0 + ): + break + + f_start[i[n], j[n]] = c + f_start[i[n] - vi[n], j[n]] = c + f_start[i[n], j[n] - vj[n]] = c + + f_end = f_start.clone() - assert ( - f_end[i[n], j[n]] == 0 - and f_end[i[n] - vi[n], j[n]] == 0 - and f_end[i[n], j[n] - vj[n]] == 0 - ) - - if (i[n] == 0 and vi[n] == -1) or (i[n] == height - 1 and vi[n] == 1): - vi[n] = -vi[n] - if (j[n] == 0 and vj[n] == -1) or (j[n] == width - 1 and vj[n] == 1): - vj[n] = -vj[n] - - i[n] += vi[n] - j[n] += vj[n] - - if not ( - f_end[i[n], j[n]] == 0 - and f_end[i[n] - vi[n], j[n]] == 0 - and f_end[i[n], j[n] - vj[n]] == 0 - ): - i[n], j[n], vi[n], vj[n] = pi, pj, pvi, pvj - - f_end[i[n], j[n]] = c - f_end[i[n] - vi[n], j[n]] = c - f_end[i[n], j[n] - vj[n]] = c + for l in range(nb_iterations): + f_end[...] = 0 + nb_collisions = 0 + for n in range(nb_birds): + c = col[n] + + pi, pj, pvi, pvj = ( + i[n].item(), + j[n].item(), + vi[n].item(), + vj[n].item(), + ) + + if (i[n] == 0 and vi[n] == -1) or ( + i[n] == height - 1 and vi[n] == 1 + ): + vi[n] = -vi[n] + if (j[n] == 0 and vj[n] == -1) or ( + j[n] == width - 1 and vj[n] == 1 + ): + vj[n] = -vj[n] + + i[n] += vi[n] + j[n] += vj[n] + + if not ( + f_end[i[n], j[n]] == 0 + and f_end[i[n] - vi[n], j[n]] == 0 + and f_end[i[n], j[n] - vj[n]] == 0 + ): + nb_collisions += 1 + + f_end[i[n], j[n]] = c + f_end[i[n] - vi[n], j[n]] = c + f_end[i[n], j[n] - vj[n]] = c + + if nb_collisions == 0: + break pairs.append((f_start, f_end)) @@ -154,7 +150,10 @@ def generate( return torch.cat(result, dim=0) -def generate_( +###################################################################### + + +def generate_seq_( nb, height, width, @@ -303,7 +302,7 @@ if __name__ == "__main__": height, width = 6, 8 start_time = time.perf_counter() - seq = generate(nb=90, height=height, width=width) + seq = generate_seq(nb=90, height=height, width=width) delay = time.perf_counter() - start_time print(f"{seq.size(0)/delay:02f} samples/s") -- 2.39.5