X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=grids.py;h=47e586119fc052ebd6404ff8046942e4bafee196;hb=9ec709a2a08eb82dfc17ef1e24aa9a84751d63e0;hp=9462f87676687503576c030b1ac111c9a1fde0c3;hpb=b269012c29f1eeae7d51625694269da40326a69f;p=culture.git diff --git a/grids.py b/grids.py index 9462f87..47e5861 100755 --- a/grids.py +++ b/grids.py @@ -23,9 +23,9 @@ class Grids(problem.Problem): ("red", [255, 0, 0]), ("green", [0, 192, 0]), ("blue", [0, 0, 255]), - ("orange", [255, 192, 0]), + ("yellow", [255, 224, 0]), ("cyan", [0, 255, 255]), - ("violet", [255, 0, 255]), + ("violet", [224, 128, 255]), ("lightgreen", [192, 255, 192]), ("brown", [165, 42, 42]), ("lightblue", [192, 192, 255]), @@ -34,7 +34,6 @@ class Grids(problem.Problem): def __init__(self, device=torch.device("cpu")): self.colors = torch.tensor([c for _, c in self.named_colors]) - self.name2color = dict([(p[0], i) for i, p in enumerate(self.named_colors)]) self.height = 10 self.width = 10 self.device = device @@ -66,19 +65,6 @@ class Grids(problem.Problem): return x - def frame2img_(self, x, scale=15): - x = x.reshape(x.size(0), self.height, -1) - x = self.colors[x].permute(0, 3, 1, 2) - s = x.shape - x = x[:, :, :, None, :, None].expand(-1, -1, -1, scale, -1, scale) - x = x.reshape(s[0], s[1], s[2] * scale, s[3] * scale) - - x[:, :, :, torch.arange(0, x.size(3), scale)] = 0 - x[:, :, torch.arange(0, x.size(2), scale), :] = 0 - x = x[:, :, 1:, 1:] - - return x - def save_image( self, result_dir, @@ -88,6 +74,7 @@ class Grids(problem.Problem): predicted_prompts=None, predicted_answers=None, nrow=4, + margin=8, ): S = self.height * self.width As = prompts[:, 0 * (S + 1) : 0 * (S + 1) + S].view(-1, self.height, self.width) @@ -134,8 +121,6 @@ class Grids(problem.Problem): return y - margin = 8 - img_prompts = torch.cat( [ add_frame( @@ -209,10 +194,39 @@ class Grids(problem.Problem): def nb_token_values(self): return len(self.colors) + # @torch.compile + def rec_coo_(self, nb_rec, min_height=3, min_width=3): + # @torch.compile + def overlap(ia, ja, ib, jb): + return ( + ia[1] >= ib[0] and ia[0] <= ib[1] and ja[1] >= jb[0] and ja[0] <= jb[1] + ) + + if nb_rec == 3: + while True: + i = torch.randint(self.height + 1, (nb_rec, 2)).sort(dim=1).values + j = torch.randint(self.width + 1, (nb_rec, 2)).sort(dim=1).values + if ( + not ( + overlap(i[0], j[0], i[1], j[1]) + or overlap(i[0], j[0], i[2], j[2]) + or overlap(i[1], j[1], i[2], j[2]) + ) + and (i[:, 1] - i[:, 0]).min() >= min_height + and (j[:, 1] - j[:, 0]).min() >= min_width + ): + break + return ( + (i[0, 0], j[0, 0], i[0, 1], j[0, 1]), + (i[1, 0], j[1, 0], i[1, 1], j[1, 1]), + (i[2, 0], j[2, 0], i[2, 1], j[2, 1]), + ) + # That's quite a tensorial spaghetti mess to sample # non-overlapping rectangles quickly, but made the generation of # 100k samples go from 1h50 with a lame pure python code to 3min30s # with this one. + # @torch.compile def rec_coo(self, nb_rec, min_height=3, min_width=3): nb_trials = 200 @@ -274,6 +288,7 @@ class Grids(problem.Problem): ) ] + # @torch.compile def rec_coo_(self, x, n, min_height=3, min_width=3): collision = x.new(x.size()) while True: @@ -298,6 +313,7 @@ class Grids(problem.Problem): ###################################################################### + # @torch.compile def task_replace_color(self, A, f_A, B, f_B): nb_rec = 3 c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1 @@ -308,6 +324,7 @@ class Grids(problem.Problem): X[i1:i2, j1:j2] = c[n] f_X[i1:i2, j1:j2] = c[n if n > 0 else -1] + # @torch.compile def task_translate(self, A, f_A, B, f_B): di, dj = torch.randint(3, (2,)) - 1 nb_rec = 3 @@ -332,6 +349,7 @@ class Grids(problem.Problem): else: f_X[i1:i2, j1:j2] = c[n] + # @torch.compile def task_grow(self, A, f_A, B, f_B): di, dj = torch.randint(2, (2,)) * 2 - 1 nb_rec = 3 @@ -357,6 +375,7 @@ class Grids(problem.Problem): X[i1:i2, j1:j2] = c[n] f_X[i1:i2, j1:j2] = c[n] + # @torch.compile def task_color_grow(self, A, f_A, B, f_B): di, dj = torch.randint(2, (2,)) * 2 - 1 nb_rec = 3 @@ -398,6 +417,7 @@ class Grids(problem.Problem): else: f_X[i1:i2, j : j + 1] = c[2 * n + 1] + # @torch.compile def task_frame(self, A, f_A, B, f_B): nb_rec = 3 c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1 @@ -410,6 +430,7 @@ class Grids(problem.Problem): if n == nb_rec - 1: f_X[i1 + 1 : i2 - 1, j1 + 1 : j2 - 1] = 0 + # @torch.compile def task_detect(self, A, f_A, B, f_B): nb_rec = 3 c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1 @@ -421,6 +442,7 @@ class Grids(problem.Problem): if n < nb_rec - 1: f_X[i1, j1] = c[-1] + # @torch.compile def contact(self, X, i, j, q): nq, nq_diag = 0, 0 no = 0 @@ -456,8 +478,9 @@ class Grids(problem.Problem): return no, nq, nq_diag + # @torch.compile def task_count(self, A, f_A, B, f_B): - N = torch.randint(4, (1,)) + 2 + N = (torch.randint(4, (1,)) + 2).item() c = torch.randperm(len(self.colors) - 1)[:N] + 1 for X, f_X in [(A, f_A), (B, f_B)]: @@ -479,6 +502,7 @@ class Grids(problem.Problem): for j in range(nb[n]): f_X[n, j] = c[n] + # @torch.compile def task_trajectory(self, A, f_A, B, f_B): c = torch.randperm(len(self.colors) - 1)[:2] + 1 for X, f_X in [(A, f_A), (B, f_B)]: @@ -506,10 +530,11 @@ class Grids(problem.Problem): f_X[i + k * di, j + k * dj] = c[min(k, 1)] k += 1 + # @torch.compile def task_bounce(self, A, f_A, B, f_B): c = torch.randperm(len(self.colors) - 1)[:3] + 1 for X, f_X in [(A, f_A), (B, f_B)]: - + # @torch.compile def free(i, j): return ( i >= 0 @@ -569,6 +594,7 @@ class Grids(problem.Problem): if l > 3: break + # @torch.compile def task_scale(self, A, f_A, B, f_B): c = torch.randperm(len(self.colors) - 1)[:2] + 1 @@ -593,38 +619,109 @@ class Grids(problem.Problem): X[i, j] = c[1] f_X[0:2, 0:2] = c[1] - def task_islands(self, A, f_A, B, f_B): + # @torch.compile + def task_symbols(self, A, f_A, B, f_B): + nb_rec = 4 + c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1 + delta = 3 for X, f_X in [(A, f_A), (B, f_B)]: while True: - i, j = torch.randint(self.height, (1,)), torch.randint(self.width, (1,)) - if ( - i == 0 - or i == self.height - 1 - or j == 0 - or j == self.width - 1 - or X[i, j] == 1 - ): - break - while True: - di, dj = torch.randint(3, (2,)) - 1 - if abs(di) + abs(dj) > 0: + i, j = torch.randint(self.height - delta + 1, (nb_rec,)), torch.randint( + self.width - delta + 1, (nb_rec,) + ) + d = (i[None, :] - i[:, None]).abs().max((j[None, :] - j[:, None]).abs()) + d.fill_diagonal_(delta + 1) + if d.min() > delta: break - X[i, j] = 1 + + for k in range(1, nb_rec): + X[i[k] : i[k] + delta, j[k] : j[k] + delta] = c[k] + + ai, aj = i.float().mean(), j.float().mean() + + q = torch.randint(3, (1,)) + 1 + + X[i[0] + delta // 2 - 1, j[0] + delta // 2 - 1] = c[0] + X[i[0] + delta // 2 - 1, j[0] + delta // 2 + 1] = c[0] + X[i[0] + delta // 2 + 1, j[0] + delta // 2 - 1] = c[0] + X[i[0] + delta // 2 + 1, j[0] + delta // 2 + 1] = c[0] + + assert i[q] != ai and j[q] != aj + + X[ + i[0] + delta // 2 + (i[q] - ai).sign().long(), + j[0] + delta // 2 + (j[q] - aj).sign().long(), + ] = c[nb_rec] + + f_X[i[0] : i[0] + delta, j[0] : j[0] + delta] = c[q] + + # @torch.compile + def task_ortho(self, A, f_A, B, f_B): + nb_rec = 3 + di, dj = torch.randint(3, (2,)) - 1 + o = torch.tensor([[0.0, 1.0], [-1.0, 0.0]]) + m = torch.eye(2) + for _ in range(torch.randint(4, (1,))): + m = m @ o + if torch.rand(1) < 0.5: + m[0, :] = -m[0, :] + + ci, cj = (self.height - 1) / 2, (self.width - 1) / 2 + + for X, f_X in [(A, f_A), (B, f_B)]: while True: - i, j = i + di, j + dj - if i < 0 or i >= self.height or j < 0 or j >= self.width: - break - b = ( - i == 0 - or i == self.height - 1 - or j == 0 - or j == self.width - 1 - or X[i, j] == 1 - ) - X[i, j] = 1 - if b: + X[...] = 0 + f_X[...] = 0 + + c = torch.randperm(len(self.colors) - 1)[:nb_rec] + 1 + + for r in range(nb_rec): + while True: + i1, i2 = torch.randint(self.height - 2, (2,)) + 1 + j1, j2 = torch.randint(self.width - 2, (2,)) + 1 + if ( + i2 >= i1 + and j2 >= j1 + and max(i2 - i1, j2 - j1) >= 2 + and min(i2 - i1, j2 - j1) <= 3 + ): + break + X[i1 : i2 + 1, j1 : j2 + 1] = c[r] + + i1, j1, i2, j2 = i1 - ci, j1 - cj, i2 - ci, j2 - cj + + i1, j1 = m[0, 0] * i1 + m[0, 1] * j1, m[1, 0] * i1 + m[1, 1] * j1 + i2, j2 = m[0, 0] * i2 + m[0, 1] * j2, m[1, 0] * i2 + m[1, 1] * j2 + + i1, j1, i2, j2 = i1 + ci, j1 + cj, i2 + ci, j2 + cj + i1, i2 = i1.long() + di, i2.long() + di + j1, j2 = j1.long() + dj, j2.long() + dj + if i1 > i2: + i1, i2 = i2, i1 + if j1 > j2: + j1, j2 = j2, j1 + + f_X[i1 : i2 + 1, j1 : j2 + 1] = c[r] + + n = F.one_hot(X.flatten()).sum(dim=0)[1:] + if ( + n.sum() > self.height * self.width // 4 + and (n > 0).long().sum() == nb_rec + ): break + # @torch.compile + def task_islands(self, A, f_A, B, f_B): + pass + + # for X, f_X in [(A, f_A), (B, f_B)]: + # n = torch.arange(self.height * self.width).reshape(self.height, self.width) + # k = torch.randperm(self.height * self.width) + # X[...]=-1 + # for q in k: + # i,j=q%self.height,q//self.height + # if + ###################################################################### def all_tasks(self): @@ -639,10 +736,20 @@ class Grids(problem.Problem): self.task_trajectory, self.task_bounce, self.task_scale, + self.task_symbols, + self.task_ortho, # self.task_islands, ] - def generate_prompts_and_answers(self, nb, tasks=None, device="cpu"): + def trivial_prompts_and_answers(self, prompts, answers): + S = self.height * self.width + Bs = prompts[:, 2 * (S + 1) : 2 * (S + 1) + S] + f_Bs = answers + return (Bs == f_Bs).long().min(dim=-1).values > 0 + + def generate_prompts_and_answers( + self, nb, tasks=None, progress_bar=False, device="cpu" + ): if tasks is None: tasks = self.all_tasks() @@ -650,12 +757,17 @@ class Grids(problem.Problem): prompts = torch.zeros(nb, 3 * S + 2, dtype=torch.int64) answers = torch.zeros(nb, S, dtype=torch.int64) - for prompt, answer in tqdm.tqdm( - zip(prompts, answers), - dynamic_ncols=True, - desc="world generation", - total=prompts.size(0), - ): + bunch = zip(prompts, answers) + + if progress_bar: + bunch = tqdm.tqdm( + bunch, + dynamic_ncols=True, + desc="world generation", + total=prompts.size(0), + ) + + for prompt, answer in bunch: A = prompt[0 * (S + 1) : 0 * (S + 1) + S].view(self.height, self.width) f_A = prompt[1 * (S + 1) : 1 * (S + 1) + S].view(self.height, self.width) B = prompt[2 * (S + 1) : 2 * (S + 1) + S].view(self.height, self.width) @@ -691,25 +803,40 @@ class Grids(problem.Problem): if __name__ == "__main__": import time - nb = 48 - grids = Grids() + # nb = 1000 + # grids = problem.MultiThreadProblem( + # grids, max_nb_cached_chunks=50, chunk_size=100, nb_threads=1 + # ) + # time.sleep(10) + # start_time = time.perf_counter() + # prompts, answers = grids.generate_prompts_and_answers(nb) + # delay = time.perf_counter() - start_time + # print(f"{prompts.size(0)/delay:02f} seq/s") + # exit(0) + + if True: + nb = 72 + + for t in grids.all_tasks(): + # for t in [grids.task_ortho]: + print(t.__name__) + prompts, answers = grids.generate_prompts_and_answers(nb, tasks=[t]) + grids.save_quizzes("/tmp", t.__name__, prompts[:nb], answers[:nb], nrow=4) + + exit(0) + + nb = 500 + for t in grids.all_tasks(): - # for t in [grids.task_islands]: - print(t.__name__) + start_time = time.perf_counter() prompts, answers = grids.generate_prompts_and_answers(nb, tasks=[t]) - grids.save_quizzes("/tmp", t.__name__, prompts[:nb], answers[:nb], nrow=4) + delay = time.perf_counter() - start_time + print(f"{t.__name__} {prompts.size(0)/delay:02f} seq/s") exit(0) - nb = 72 - - start_time = time.perf_counter() - prompts, answers = grids.generate_prompts_and_answers(nb) - delay = time.perf_counter() - start_time - print(f"{prompts.size(0)/delay:02f} seq/s") - m = torch.randint(2, (prompts.size(0),)) predicted_prompts = m * (torch.randint(2, (prompts.size(0),)) * 2 - 1) predicted_answers = (1 - m) * (torch.randint(2, (prompts.size(0),)) * 2 - 1)