From b76e3f632315c63dbd8f11a53b187f23057e4e1f Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Tue, 25 Jun 2024 17:43:38 +0200 Subject: [PATCH] Update. --- quizz_machine.py | 28 +++---- sky.py | 192 ++++++++++++++++++++++++++--------------------- 2 files changed, 115 insertions(+), 105 deletions(-) diff --git a/quizz_machine.py b/quizz_machine.py index 43fd868..a3da365 100755 --- a/quizz_machine.py +++ b/quizz_machine.py @@ -71,7 +71,7 @@ import sky class QuizzMachine: def save_image(self, input, result_dir, filename, logger): - img = sky.seq2img(input.to("cpu"), self.height, self.width) + img = self.sky.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) logger(f"wrote {image_name}") @@ -94,18 +94,12 @@ class QuizzMachine: ): super().__init__() + self.sky = sky.Sky(height=6, width=8, nb_birds=3, nb_iterations=2) self.batch_size = batch_size self.device = device - self.height = 6 - self.width = 8 - self.train_w_quizzes = sky.generate_seq( - nb_train_samples, height=self.height, width=self.width - ).to(device) - - self.test_w_quizzes = sky.generate_seq( - nb_test_samples, height=self.height, width=self.width - ).to(device) + self.train_w_quizzes = self.sky.generate_seq(nb_train_samples).to(device) + self.test_w_quizzes = self.sky.generate_seq(nb_test_samples).to(device) self.nb_codes = max(self.train_w_quizzes.max(), self.test_w_quizzes.max()) + 1 @@ -234,9 +228,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:] = sky.generate_seq(nb, height=self.height, width=self.width).to( - self.device - ) + input[-nb:] = self.sky.generate_seq(nb).to(self.device) def store_c_quizzes(self, new_c_quizzes, for_train=True): if for_train: @@ -258,7 +250,7 @@ class QuizzMachine: # Generate quizzes with model c_quizzes = torch.empty( - nb, self.height * self.width * 2 + 1, device=self.device, dtype=torch.int64 + nb, self.train_w_quizzes.size(1), device=self.device, dtype=torch.int64 ) ar_mask = torch.full(c_quizzes.size(), 1, device=self.device) @@ -306,11 +298,11 @@ class QuizzMachine: ############################################################### # Create the reverse quizzes - l = self.height * self.width + l = (c_quizzes.size(1) - 1) // 2 direction = c_quizzes[:, l : l + 1] - direction = sky.token_forward * ( - direction == sky.token_backward - ) + sky.token_backward * (direction == sky.token_forward) + direction = self.sky.token_forward * ( + direction == self.sky.token_backward + ) + self.sky.token_backward * (direction == self.sky.token_forward) reverse_c_quizzes = torch.cat( [c_quizzes[:, l + 1 :], direction, c_quizzes[:, :l]], dim=1 ) diff --git a/sky.py b/sky.py index 3458d85..a90e37d 100755 --- a/sky.py +++ b/sky.py @@ -15,39 +15,40 @@ from torch.nn import functional as F ###################################################################### -colors = torch.tensor( - [ - [255, 255, 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], - ] -) - -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)]) + "><" +class Sky: + colors = torch.tensor( + [ + [255, 255, 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], + ] + ) + 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 -class Sky: - def __init__(self, height, width): - self.heigh = heigh + token2char = ( + "_" + "".join([chr(ord("A") + n) for n in range(len(colors) - 1)]) + "><" + ) + + def __init__(self, height=6, width=8, nb_birds=3, nb_iterations=2): + self.height = height self.width = width + self.nb_birds = nb_birds + self.nb_iterations = nb_iterations - def generate_seq( - nb, height, width, nb_birds=3, nb_iterations=2, return_iterations=False - ): + def generate_seq(self, nb, return_iterations=False): pairs = [] kept_iterations = [] @@ -55,32 +56,37 @@ class Sky: while True: iterations = [] - f_start = torch.zeros(height, width, dtype=torch.int64) + f_start = torch.zeros(self.height, self.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), + torch.empty(self.nb_birds, dtype=torch.int64), + torch.empty(self.nb_birds, dtype=torch.int64), + torch.empty(self.nb_birds, dtype=torch.int64), + torch.empty(self.nb_birds, dtype=torch.int64), ) - col = torch.randperm(colors.size(0) - 1)[:nb_birds].sort().values + 1 + col = ( + torch.randperm(self.colors.size(0) - 1)[: self.nb_birds] + .sort() + .values + + 1 + ) - for n in range(nb_birds): + for n in range(self.nb_birds): c = col[n] while True: i[n], j[n] = ( - torch.randint(height, (1,))[0], - torch.randint(width, (1,))[0], + torch.randint(self.height, (1,))[0], + torch.randint(self.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 i[n] - vi[n] < self.height and j[n] - vj[n] >= 0 - and j[n] - vj[n] < width + and j[n] - vj[n] < self.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 @@ -93,11 +99,11 @@ class Sky: f_end = f_start.clone() - for l in range(nb_iterations): + for l in range(self.nb_iterations): iterations.append(f_end.clone()) f_end[...] = 0 nb_collisions = 0 - for n in range(nb_birds): + for n in range(self.nb_birds): c = col[n] pi, pj, pvi, pvj = ( @@ -108,11 +114,11 @@ class Sky: ) if (i[n] == 0 and vi[n] == -1) or ( - i[n] == height - 1 and vi[n] == 1 + i[n] == self.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 + j[n] == self.width - 1 and vj[n] == 1 ): vj[n] = -vj[n] @@ -143,7 +149,11 @@ class Sky: if torch.rand(1) < 0.5: result.append( torch.cat( - [p[0].flatten(), torch.tensor([token_forward]), p[1].flatten()], + [ + p[0].flatten(), + torch.tensor([self.token_forward]), + p[1].flatten(), + ], dim=0, )[None, :] ) @@ -152,7 +162,7 @@ class Sky: torch.cat( [ p[1].flatten(), - torch.tensor([token_backward]), + torch.tensor([self.token_backward]), p[0].flatten(), ], dim=0, @@ -168,27 +178,26 @@ class Sky: ###################################################################### def generate_seq_old( + self, nb, - height, - width, - nb_birds=3, - nb_iterations=2, ): pairs = [] for n in tqdm.tqdm(range(nb), dynamic_ncols=True, desc="world generation"): - f_start = torch.zeros(height, width, dtype=torch.int64) - f_end = torch.zeros(height, width, dtype=torch.int64) + f_start = torch.zeros(self.height, self.width, dtype=torch.int64) + f_end = torch.zeros(self.height, self.width, dtype=torch.int64) n = torch.arange(f_start.size(0)) for c in ( - (torch.randperm(nb_bird_tokens) + first_bird_token)[:nb_birds] + (torch.randperm(self.nb_bird_tokens) + self.first_bird_token)[ + : self.nb_birds + ] .sort() .values ): i, j = ( - torch.randint(height - 2, (1,))[0] + 1, - torch.randint(width - 2, (1,))[0] + 1, + torch.randint(self.height - 2, (1,))[0] + 1, + torch.randint(self.width - 2, (1,))[0] + 1, ) vm = torch.randint(4, (1,))[0] vi, vj = (vm // 2) * (2 * (vm % 2) - 1), (1 - vm // 2) * ( @@ -200,10 +209,10 @@ class Sky: f_start[i + vj, j - vi] = c f_start[i - vj, j + vi] = c - for l in range(nb_iterations): + for l in range(self.nb_iterations): i += vi j += vj - if i < 0 or i >= height or j < 0 or j >= width: + if i < 0 or i >= self.height or j < 0 or j >= self.width: i -= vi j -= vj vi, vj = -vi, -vj @@ -222,7 +231,11 @@ class Sky: if torch.rand(1) < 0.5: result.append( torch.cat( - [p[0].flatten(), torch.tensor([token_forward]), p[1].flatten()], + [ + p[0].flatten(), + torch.tensor([self.token_forward]), + p[1].flatten(), + ], dim=0, )[None, :] ) @@ -231,7 +244,7 @@ class Sky: torch.cat( [ p[1].flatten(), - torch.tensor([token_backward]), + torch.tensor([self.token_backward]), p[0].flatten(), ], dim=0, @@ -240,10 +253,12 @@ class Sky: return torch.cat(result, dim=0) - def frame2img(x, height, width, upscale=15): - x = x.reshape(-1, height, width) - m = torch.logical_and(x >= 0, x < first_bird_token + nb_bird_tokens).long() - x = colors[x * m].permute(0, 3, 1, 2) + def frame2img(self, x, upscale=15): + x = x.reshape(-1, self.height, self.width) + m = torch.logical_and( + x >= 0, x < self.first_bird_token + self.nb_bird_tokens + ).long() + x = self.colors[x * m].permute(0, 3, 1, 2) s = x.shape x = x[:, :, :, None, :, None].expand(-1, -1, -1, upscale, -1, upscale) x = x.reshape(s[0], s[1], s[2] * upscale, s[3] * upscale) @@ -262,61 +277,65 @@ class Sky: return x - def seq2img(seq, height, width, upscale=15): - f_first = seq[:, : height * width].reshape(-1, height, width) - f_second = seq[:, height * width + 1 :].reshape(-1, height, width) - direction = seq[:, height * width] + def seq2img(self, seq, upscale=15): + f_first = seq[:, : self.height * self.width].reshape( + -1, self.height, self.width + ) + f_second = seq[:, self.height * self.width + 1 :].reshape( + -1, self.height, self.width + ) + direction = seq[:, self.height * self.width] direction_symbol = torch.full( - (direction.size(0), height * upscale - 1, upscale), 0 + (direction.size(0), self.height * upscale - 1, upscale), 0 ) - direction_symbol = colors[direction_symbol].permute(0, 3, 1, 2) - separator = torch.full((direction.size(0), 3, height * upscale - 1, 1), 0) + direction_symbol = self.colors[direction_symbol].permute(0, 3, 1, 2) + separator = torch.full((direction.size(0), 3, self.height * upscale - 1, 1), 0) for n in range(direction_symbol.size(0)): - if direction[n] == token_forward: + if direction[n] == self.token_forward: for k in range(upscale): direction_symbol[ n, :, - (height * upscale) // 2 - upscale // 2 + k, + (self.height * upscale) // 2 - upscale // 2 + k, 3 + upscale // 2 - abs(k - upscale // 2), ] = 0 - elif direction[n] == token_backward: + elif direction[n] == self.token_backward: for k in range(upscale): direction_symbol[ n, :, - (height * upscale) // 2 - upscale // 2 + k, + (self.height * upscale) // 2 - upscale // 2 + k, 3 + abs(k - upscale // 2), ] = 0 else: for k in range(2, upscale - 2): direction_symbol[ - n, :, (height * upscale) // 2 - upscale // 2 + k, k + n, :, (self.height * upscale) // 2 - upscale // 2 + k, k ] = 0 direction_symbol[ n, :, - (height * upscale) // 2 - upscale // 2 + k, + (self.height * upscale) // 2 - upscale // 2 + k, upscale - 1 - k, ] = 0 return torch.cat( [ - frame2img(f_first, height, width, upscale), + self.frame2img(f_first, upscale), separator, direction_symbol, separator, - frame2img(f_second, height, width, upscale), + self.frame2img(f_second, upscale), ], dim=3, ) - def seq2str(seq): + def seq2str(self, seq): result = [] for s in seq: - result.append("".join([token2char[v] for v in s])) + result.append("".join([self.token2char[v] for v in s])) return result @@ -325,18 +344,17 @@ class Sky: if __name__ == "__main__": import time - height, width = 6, 8 + sky = Sky(height=6, width=8, nb_iterations=100) + start_time = time.perf_counter() - seq, it = generate_seq( - nb=64, height=height, width=width, nb_iterations=100, return_iterations=True - ) + seq, it = sky.generate_seq(nb=64, return_iterations=True) delay = time.perf_counter() - start_time print(f"{seq.size(0)/delay:02f} samples/s") - print(seq2str(seq[:4])) + print(sky.seq2str(seq[:4])) for t in range(len(it[0])): - img = torch.cat([frame2img(f[t], height, width) for f in it], dim=0) + img = torch.cat([sky.frame2img(f[t]) for f in it], dim=0) torchvision.utils.save_image( img.float() / 255.0, f"/tmp/frame_{t:03d}.png", @@ -348,7 +366,7 @@ if __name__ == "__main__": # m = (torch.rand(seq.size()) < 0.05).long() # seq = (1 - m) * seq + m * 23 - img = seq2img(seq, height, width) + img = sky.seq2img(seq) print(img.size()) torchvision.utils.save_image( -- 2.39.5