3 # Any copyright is dedicated to the Public Domain.
4 # https://creativecommons.org/publicdomain/zero/1.0/
6 # Written by Francois Fleuret <francois@fleuret.org>
8 import math, sys, tqdm, os, warnings
10 import torch, torchvision
13 from torch.nn import functional as F
15 ######################################################################
20 def grow_islands(nb, height, width, nb_seeds, nb_iterations):
21 w = torch.empty(5, 1, 3, 3)
23 w[0, 0] = torch.tensor(
31 w[1, 0] = torch.tensor(
39 w[2, 0] = torch.tensor(
47 w[3, 0] = torch.tensor(
55 w[4, 0] = torch.tensor(
63 Z = torch.zeros(nb, height, width)
66 for _ in range(nb_seeds):
67 M = F.conv2d(Z[:, None, :, :], w, padding=1)
68 M = torch.cat([M[:, :1], M[:, 1:].min(dim=1, keepdim=True).values], dim=1)
69 M = ((M[:, 0] == 0) & (Z == 0)).long()
70 Q = (M.flatten(1).max(dim=1).values > 0).long()[:, None]
71 M = M * torch.rand(M.size())
73 M = F.one_hot(M.argmax(dim=1), num_classes=M.size(1))
76 for _ in range(nb_iterations):
77 M = F.conv2d(Z[:, None, :, :], w, padding=1)
78 M = torch.cat([M[:, :1], M[:, 1:].min(dim=1, keepdim=True).values], dim=1)
79 M = ((M[:, 1] >= 0) & (Z == 0)).long()
80 Q = (M.flatten(1).max(dim=1).values > 0).long()[:, None]
81 M = M * torch.rand(M.size())
83 M = F.one_hot(M.argmax(dim=1), num_classes=M.size(1))
88 Z = Z * (torch.arange(Z.size(1) * Z.size(2)) + 1).reshape(1, Z.size(1), Z.size(2))
92 Z = F.max_pool2d(Z, 3, 1, 1) * M
98 V = F.one_hot(U).max(dim=1).values
99 W = V.cumsum(dim=1) - V
100 N = torch.arange(Z.size(0))[:, None, None].expand_as(Z)
106 class Grids(problem.Problem):
108 ("white", [255, 255, 255]),
109 ("red", [255, 0, 0]),
110 ("green", [0, 192, 0]),
111 ("blue", [0, 0, 255]),
112 ("yellow", [255, 224, 0]),
113 ("cyan", [0, 255, 255]),
114 ("violet", [224, 128, 255]),
115 ("lightgreen", [192, 255, 192]),
116 ("brown", [165, 42, 42]),
117 ("lightblue", [192, 192, 255]),
118 ("gray", [128, 128, 128]),
123 max_nb_cached_chunks=None,
128 self.colors = torch.tensor([c for _, c in self.named_colors])
131 self.cache_rec_coo = {}
134 self.task_replace_color,
141 self.task_trajectory,
150 self.all_tasks = all_tasks
152 self.all_tasks = [getattr(self, "task_" + t) for t in tasks.split(",")]
154 super().__init__(max_nb_cached_chunks, chunk_size, nb_threads)
156 ######################################################################
158 def frame2img(self, x, scale=15):
159 x = x.reshape(x.size(0), self.height, -1)
160 m = torch.logical_and(x >= 0, x < self.nb_token_values()).long()
161 x = self.colors[x * m].permute(0, 3, 1, 2)
163 x = x[:, :, :, None, :, None].expand(-1, -1, -1, scale, -1, scale)
164 x = x.reshape(s[0], s[1], s[2] * scale, s[3] * scale)
166 x[:, :, :, torch.arange(0, x.size(3), scale)] = 0
167 x[:, :, torch.arange(0, x.size(2), scale), :] = 0
170 for n in range(m.size(0)):
171 for i in range(m.size(1)):
172 for j in range(m.size(2)):
174 for k in range(2, scale - 2):
176 x[n, :, i * scale + k, j * scale + k - l] = 0
178 n, :, i * scale + scale - 1 - k, j * scale + k - l
189 predicted_prompts=None,
190 predicted_answers=None,
194 S = self.height * self.width
195 As = prompts[:, 0 * (S + 1) : 0 * (S + 1) + S].view(-1, self.height, self.width)
196 f_As = prompts[:, 1 * (S + 1) : 1 * (S + 1) + S].view(
197 -1, self.height, self.width
199 Bs = prompts[:, 2 * (S + 1) : 2 * (S + 1) + S].view(-1, self.height, self.width)
200 prompts = torch.cat([As, f_As, Bs], dim=2)
201 answers = answers.reshape(answers.size(0), self.height, self.width)
203 if predicted_prompts is None:
204 predicted_prompts = 255
206 if predicted_answers is None:
207 predicted_answers = 255
209 def add_frame(x, c, margin, bottom=False):
211 h, w, di, dj = x.size(2) + margin, x.size(3), 0, 0
214 x.size(2) + 2 * margin,
215 x.size(3) + 2 * margin,
220 y = x.new_full((x.size(0), x.size(1), h, w), 0)
225 c = c.long()[:, None]
227 (1 - ((c == 1).long() + (c == 0).long() + (c == -1).long()))
228 * torch.tensor([64, 64, 64])
229 + (c == 1).long() * torch.tensor([0, 255, 0])
230 + (c == 0).long() * torch.tensor([255, 255, 255])
231 + (c == -1).long() * torch.tensor([255, 0, 0])
233 y[...] = c[:, :, None, None]
235 y[:, :, di : di + x.size(2), dj : dj + x.size(3)] = x
239 img_prompts = torch.cat(
242 add_frame(self.frame2img(x), c=0, margin=1),
246 for x in prompts.to("cpu").split(split_size=self.width, dim=2)
251 h = img_prompts.size(2)
252 img_answers = add_frame(
253 add_frame(self.frame2img(answers.to("cpu")), c=0, margin=1),
258 separator_size = 2 * margin
260 separator = img_prompts.new_full(
270 marker = img_prompts.new_full(
280 # marker[:, :, 0] = 0
281 # marker[:, :, h - 1] = 0
283 for k in range(1, 2 * separator_size - 8):
284 i = k - (separator_size - 4)
285 j = separator_size - 5 - abs(i)
286 marker[:, :, h // 2 - 1 + i, 2 + j] = 0
287 marker[:, :, h // 2 - 1 + i + 1, 2 + j] = 0
298 image_name = os.path.join(result_dir, filename)
299 torchvision.utils.save_image(
307 ######################################################################
309 def nb_token_values(self):
310 return len(self.colors)
319 prevent_overlap=False,
321 if surface_max is None:
322 surface_max = self.height * self.width // 2
324 signature = (nb_rec, min_height, min_width, surface_max)
327 return self.cache_rec_coo[signature].pop()
336 i = torch.randint(self.height, (N * nb_rec, 2)).sort(dim=-1).values
337 j = torch.randint(self.width, (N * nb_rec, 2)).sort(dim=-1).values
340 (i[:, 1] >= i[:, 0] + min_height)
341 & (j[:, 1] >= j[:, 0] + min_height)
342 & ((i[:, 1] - i[:, 0]) * (j[:, 1] - j[:, 0]) <= surface_max)
345 i, j = i[big_enough], j[big_enough]
347 n = i.size(0) - i.size(0) % nb_rec
352 i = i[:n].reshape(n // nb_rec, nb_rec, -1)
353 j = j[:n].reshape(n // nb_rec, nb_rec, -1)
356 can_fit = ((i[:, :, 1] - i[:, :, 0]) * (j[:, :, 1] - j[:, :, 0])).sum(
358 ) <= self.height * self.width
359 i, j = i[can_fit], j[can_fit]
361 A_i1, A_i2, A_j1, A_j2 = (
367 B_i1, B_i2, B_j1, B_j2 = (
373 no_overlap = torch.logical_not(
379 i, j = i[no_overlap], j[no_overlap]
381 A_i1, A_i2, A_j1, A_j2 = (
387 B_i1, B_i2, B_j1, B_j2 = (
393 C_i1, C_i2, C_j1, C_j2 = (
419 i, j = (i[no_overlap], j[no_overlap])
426 self.cache_rec_coo[signature] = [
434 for k in range(nb_rec)
436 for n in range(i.size(0))
439 return self.cache_rec_coo[signature].pop()
441 ######################################################################
444 def task_replace_color(self, A, f_A, B, f_B):
446 c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
447 for X, f_X in [(A, f_A), (B, f_B)]:
448 r = self.rec_coo(nb_rec, prevent_overlap=True)
449 for n in range(nb_rec):
450 i1, j1, i2, j2 = r[n]
451 X[i1:i2, j1:j2] = c[n]
452 f_X[i1:i2, j1:j2] = c[n if n > 0 else -1]
455 def task_translate(self, A, f_A, B, f_B):
457 di, dj = torch.randint(3, (2,)) - 1
458 if di.abs() + dj.abs() > 0:
462 c = torch.randperm(len(self.colors) - 1)[:nb_rec] + 1
463 for X, f_X in [(A, f_A), (B, f_B)]:
465 r = self.rec_coo(nb_rec, prevent_overlap=True)
466 i1, j1, i2, j2 = r[nb_rec - 1]
469 and i2 + di < X.size(0)
471 and j2 + dj < X.size(1)
475 for n in range(nb_rec):
476 i1, j1, i2, j2 = r[n]
477 X[i1:i2, j1:j2] = c[n]
479 f_X[i1 + di : i2 + di, j1 + dj : j2 + dj] = c[n]
481 f_X[i1:i2, j1:j2] = c[n]
484 def task_grow(self, A, f_A, B, f_B):
485 di, dj = torch.randint(2, (2,)) * 2 - 1
487 c = torch.randperm(len(self.colors) - 1)[:nb_rec] + 1
488 direction = torch.randint(2, (1,)).item()
489 for X, f_X in [(A, f_A), (B, f_B)]:
491 r = self.rec_coo(nb_rec, prevent_overlap=True)
492 i1, j1, i2, j2 = r[nb_rec - 1]
493 if i1 + 3 < i2 and j1 + 3 < j2:
496 for n in range(nb_rec):
497 i1, j1, i2, j2 = r[n]
500 X[i1 + 1 : i2 - 1, j1 + 1 : j2 - 1] = c[n]
501 f_X[i1:i2, j1:j2] = c[n]
503 X[i1:i2, j1:j2] = c[n]
504 f_X[i1 + 1 : i2 - 1, j1 + 1 : j2 - 1] = c[n]
506 X[i1:i2, j1:j2] = c[n]
507 f_X[i1:i2, j1:j2] = c[n]
510 def task_half_fill(self, A, f_A, B, f_B):
511 di, dj = torch.randint(2, (2,)) * 2 - 1
513 c = torch.randperm(len(self.colors) - 1)[: 2 * nb_rec] + 1
514 direction = torch.randint(4, (1,)).item()
515 for X, f_X in [(A, f_A), (B, f_B)]:
516 r = self.rec_coo(nb_rec, prevent_overlap=True)
517 for n in range(nb_rec):
518 i1, j1, i2, j2 = r[n]
519 X[i1:i2, j1:j2] = c[2 * n]
520 f_X[i1:i2, j1:j2] = c[2 * n]
521 # Not my proudest moment
524 X[i : i + 1, j1:j2] = c[2 * n + 1]
526 f_X[i:i2, j1:j2] = c[2 * n + 1]
528 f_X[i : i + 1, j1:j2] = c[2 * n + 1]
530 i = (i1 + i2 - 1) // 2
531 X[i : i + 1, j1:j2] = c[2 * n + 1]
533 f_X[i1 : i + 1, j1:j2] = c[2 * n + 1]
535 f_X[i : i + 1, j1:j2] = c[2 * n + 1]
538 X[i1:i2, j : j + 1] = c[2 * n + 1]
540 f_X[i1:i2, j:j2] = c[2 * n + 1]
542 f_X[i1:i2, j : j + 1] = c[2 * n + 1]
544 j = (j1 + j2 - 1) // 2
545 X[i1:i2, j : j + 1] = c[2 * n + 1]
547 f_X[i1:i2, j1 : j + 1] = c[2 * n + 1]
549 f_X[i1:i2, j : j + 1] = c[2 * n + 1]
552 def task_frame(self, A, f_A, B, f_B):
554 c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
555 for X, f_X in [(A, f_A), (B, f_B)]:
556 r = self.rec_coo(nb_rec, prevent_overlap=True)
557 for n in range(nb_rec):
558 i1, j1, i2, j2 = r[n]
559 X[i1:i2, j1:j2] = c[n]
561 f_X[i1:i2, j1] = c[n]
562 f_X[i1:i2, j2 - 1] = c[n]
563 f_X[i1, j1:j2] = c[n]
564 f_X[i2 - 1, j1:j2] = c[n]
566 f_X[i1:i2, j1:j2] = c[n]
569 def task_detect(self, A, f_A, B, f_B):
571 c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
572 for X, f_X in [(A, f_A), (B, f_B)]:
573 r = self.rec_coo(nb_rec, prevent_overlap=True)
574 for n in range(nb_rec):
575 i1, j1, i2, j2 = r[n]
576 X[i1:i2, j1:j2] = c[n]
581 def contact(self, X, i, j, q):
595 if ii >= 0 and ii < self.height and jj >= 0 and jj < self.width:
596 if X[ii, jj] != 0 and X[ii, jj] != q:
605 if ii >= 0 and ii < self.height and jj >= 0 and jj < self.width:
606 if X[ii, jj] == q and X[i, jj] != q and X[ii, j] != q:
609 for ii, jj in [(i - 1, j), (i, j - 1), (i, j + 1), (i + 1, j)]:
610 if ii >= 0 and ii < self.height and jj >= 0 and jj < self.width:
614 return no, nq, nq_diag
616 def task_count(self, A, f_A, B, f_B):
620 N = torch.randint(5, (1,)).item() + 1
621 c = torch.zeros(N + 1)
622 c[1:] = torch.randperm(len(self.colors) - 1)[:N] + 1
624 for X, f_X in [(A, f_A), (B, f_B)]:
625 if not hasattr(self, "cache_count") or len(self.cache_count) == 0:
626 self.cache_count = list(
631 nb_seeds=self.height * self.width // 8,
632 nb_iterations=self.height * self.width // 10,
636 X[...] = self.cache_count.pop()
638 k = (X.max() + 1 + (c.size(0) - 1)).item()
639 V = torch.arange(k) // (c.size(0) - 1)
640 V = (V + torch.rand(V.size())).sort().indices[: X.max() + 1] % (
646 if F.one_hot(X.flatten()).max(dim=0).values.sum().item() == N + 1:
648 for e in range(1, N + 1):
649 for j in range((X == c[e]).sum() + 1):
663 def task_trajectory(self, A, f_A, B, f_B):
664 c = torch.randperm(len(self.colors) - 1)[:2] + 1
665 for X, f_X in [(A, f_A), (B, f_B)]:
667 di, dj = torch.randint(7, (2,)) - 3
669 torch.randint(self.height, (1,)).item(),
670 torch.randint(self.width, (1,)).item(),
673 abs(di) + abs(dj) > 0
675 and i + 2 * di < self.height
677 and j + 2 * dj < self.width
684 and i + k * di < self.height
686 and j + k * dj < self.width
689 X[i + k * di, j + k * dj] = c[k]
690 f_X[i + k * di, j + k * dj] = c[min(k, 1)]
694 def task_bounce(self, A, f_A, B, f_B):
695 c = torch.randperm(len(self.colors) - 1)[:3] + 1
696 for X, f_X in [(A, f_A), (B, f_B)]:
711 for _ in range((self.height * self.width) // 10):
713 torch.randint(self.height, (1,)).item(),
714 torch.randint(self.width, (1,)).item(),
720 di, dj = torch.randint(7, (2,)) - 3
721 if abs(di) + abs(dj) == 1:
725 torch.randint(self.height, (1,)).item(),
726 torch.randint(self.width, (1,)).item(),
735 if free(i + di, j + dj):
737 elif free(i - dj, j + di):
739 if free(i + dj, j - di):
740 if torch.rand(1) < 0.5:
742 elif free(i + dj, j - di):
747 i, j = i + di, j + dj
762 def task_scale(self, A, f_A, B, f_B):
763 c = torch.randperm(len(self.colors) - 1)[:2] + 1
766 torch.randint(self.height // 2, (1,)).item(),
767 torch.randint(self.width // 2, (1,)).item(),
770 for X, f_X in [(A, f_A), (B, f_B)]:
774 torch.randint(self.height // 2 + 1, (1,)).item(),
775 torch.randint(self.width // 2 + 1, (1,)).item(),
778 torch.randint(self.height // 2 + 1, (1,)).item(),
779 torch.randint(self.width // 2 + 1, (1,)).item(),
781 if i1 < i2 and j1 < j2 and min(i2 - i1, j2 - j1) <= 3:
783 X[i + i1 : i + i2, j + j1 : j + j2] = c[0]
784 f_X[2 * i1 : 2 * i2, 2 * j1 : 2 * j2] = c[0]
790 def task_symbols(self, A, f_A, B, f_B):
792 c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
794 for X, f_X in [(A, f_A), (B, f_B)]:
796 i, j = torch.randint(self.height - delta + 1, (nb_rec,)), torch.randint(
797 self.width - delta + 1, (nb_rec,)
799 d = (i[None, :] - i[:, None]).abs().max((j[None, :] - j[:, None]).abs())
800 d.fill_diagonal_(delta + 1)
804 for k in range(1, nb_rec):
805 X[i[k] : i[k] + delta, j[k] : j[k] + delta] = c[k]
807 ai, aj = i.float().mean(), j.float().mean()
809 q = torch.randint(3, (1,)).item() + 1
811 X[i[0] + delta // 2 - 1, j[0] + delta // 2 - 1] = c[0]
812 X[i[0] + delta // 2 - 1, j[0] + delta // 2 + 1] = c[0]
813 X[i[0] + delta // 2 + 1, j[0] + delta // 2 - 1] = c[0]
814 X[i[0] + delta // 2 + 1, j[0] + delta // 2 + 1] = c[0]
816 assert i[q] != ai and j[q] != aj
819 i[0] + delta // 2 + (i[q] - ai).sign().long(),
820 j[0] + delta // 2 + (j[q] - aj).sign().long(),
823 f_X[i[0] : i[0] + delta, j[0] : j[0] + delta] = c[q]
826 def task_isometry(self, A, f_A, B, f_B):
828 di, dj = torch.randint(3, (2,)) - 1
829 o = torch.tensor([[0.0, 1.0], [-1.0, 0.0]])
831 for _ in range(torch.randint(4, (1,)).item()):
833 if torch.rand(1) < 0.5:
836 ci, cj = (self.height - 1) / 2, (self.width - 1) / 2
838 for X, f_X in [(A, f_A), (B, f_B)]:
843 c = torch.randperm(len(self.colors) - 1)[:nb_rec] + 1
845 for r in range(nb_rec):
847 i1, i2 = torch.randint(self.height - 2, (2,)) + 1
848 j1, j2 = torch.randint(self.width - 2, (2,)) + 1
852 and max(i2 - i1, j2 - j1) >= 2
853 and min(i2 - i1, j2 - j1) <= 3
856 X[i1 : i2 + 1, j1 : j2 + 1] = c[r]
858 i1, j1, i2, j2 = i1 - ci, j1 - cj, i2 - ci, j2 - cj
860 i1, j1 = m[0, 0] * i1 + m[0, 1] * j1, m[1, 0] * i1 + m[1, 1] * j1
861 i2, j2 = m[0, 0] * i2 + m[0, 1] * j2, m[1, 0] * i2 + m[1, 1] * j2
863 i1, j1, i2, j2 = i1 + ci, j1 + cj, i2 + ci, j2 + cj
864 i1, i2 = i1.long() + di, i2.long() + di
865 j1, j2 = j1.long() + dj, j2.long() + dj
871 f_X[i1 : i2 + 1, j1 : j2 + 1] = c[r]
873 n = F.one_hot(X.flatten()).sum(dim=0)[1:]
875 n.sum() > self.height * self.width // 4
876 and (n > 0).long().sum() == nb_rec
880 def compute_distance(self, walls, goal_i, goal_j, start_i, start_j):
881 max_length = walls.numel()
882 dist = torch.full_like(walls, max_length)
884 dist[goal_i, goal_j] = 0
885 pred_dist = torch.empty_like(dist)
888 pred_dist.copy_(dist)
892 dist[None, 1:-1, 0:-2],
893 dist[None, 2:, 1:-1],
894 dist[None, 1:-1, 2:],
895 dist[None, 0:-2, 1:-1],
902 dist[1:-1, 1:-1].minimum_(d) # = torch.min(dist[1:-1, 1:-1], d)
903 dist = walls * max_length + (1 - walls) * dist
905 if dist[start_i, start_j] < max_length or dist.equal(pred_dist):
906 return dist * (1 - walls)
909 def task_path(self, A, f_A, B, f_B):
910 c = torch.randperm(len(self.colors) - 1)[:3] + 1
911 dist = torch.empty(self.height + 2, self.width + 2)
912 for X, f_X in [(A, f_A), (B, f_B)]:
913 nb_rec = torch.randint(3, (1,)).item() + 1
915 r = self.rec_coo(nb_rec, prevent_overlap=True)
918 for n in range(nb_rec):
919 i1, j1, i2, j2 = r[n]
920 X[i1:i2, j1:j2] = c[0]
921 f_X[i1:i2, j1:j2] = c[0]
924 torch.randint(self.height, (1,)).item(),
925 torch.randint(self.width, (1,)).item(),
931 torch.randint(self.height, (1,)).item(),
932 torch.randint(self.width, (1,)).item(),
937 dist[1:-1, 1:-1] = (X != 0).long()
938 dist[...] = self.compute_distance(dist, i1 + 1, j1 + 1, i0 + 1, j0 + 1)
939 if dist[i0 + 1, j0 + 1] >= 1 and dist[i0 + 1, j0 + 1] < self.height * 4:
942 dist[1:-1, 1:-1] += (X != 0).long() * self.height * self.width
943 dist[0, :] = self.height * self.width
944 dist[-1, :] = self.height * self.width
945 dist[:, 0] = self.height * self.width
946 dist[:, -1] = self.height * self.width
947 # dist += torch.rand(dist.size())
949 i, j = i0 + 1, j0 + 1
950 while i != i1 + 1 or j != j1 + 1:
951 f_X[i - 1, j - 1] = c[2]
974 # for X, f_X in [(A, f_A), (B, f_B)]:
975 # n = torch.arange(self.height * self.width).reshape(self.height, self.width)
976 # k = torch.randperm(self.height * self.width)
979 # i,j=q%self.height,q//self.height
983 def task_puzzle(self, A, f_A, B, f_B):
985 i0, j0 = (self.height - S) // 2, (self.width - S) // 2
986 c = torch.randperm(len(self.colors) - 1)[:4] + 1
987 for X, f_X in [(A, f_A), (B, f_B)]:
990 h = list(torch.randperm(c.size(0)))
991 n = torch.zeros(c.max() + 1)
993 k = torch.randperm(S * S)
995 i, j = q % S + i0, q // S + j0
1003 r, s, t, u = torch.tensor([r, s, t, u])[torch.randperm(4)]
1004 if r > 0 and n[r] < 6:
1007 elif s > 0 and n[s] < 6:
1010 elif t > 0 and n[t] < 6:
1013 elif u > 0 and n[u] < 6:
1022 if n.sum() == S * S:
1029 torch.randint(self.height, (1,)).item(),
1030 torch.randint(self.width, (1,)).item(),
1036 ii + i >= self.height
1037 or jj + j >= self.width
1039 f_X[i + i0, j + j0] == c[d]
1040 and X[ii + i, jj + j] > 0
1048 if f_X[i + i0, j + j0] == c[d]:
1049 X[ii + i, jj + j] = c[d]
1051 def task_islands(self, A, f_A, B, f_B):
1052 c = torch.randperm(len(self.colors) - 1)[:2] + 1
1053 for X, f_X in [(A, f_A), (B, f_B)]:
1054 if not hasattr(self, "cache_islands") or len(self.cache_islands) == 0:
1055 self.cache_islands = list(
1060 nb_seeds=self.height * self.width // 20,
1061 nb_iterations=self.height * self.width // 2,
1065 A = self.cache_islands.pop()
1069 torch.randint(self.height // 2, (1,)).item(),
1070 torch.randint(self.width // 2, (1,)).item(),
1075 X[...] = (A > 0) * c[0]
1077 f_X[...] = (A == A[i, j]) * c[1] + ((A > 0) & (A != A[i, j])) * c[0]
1079 ######################################################################
1081 def trivial_prompts_and_answers(self, prompts, answers):
1082 S = self.height * self.width
1083 Bs = prompts[:, 2 * (S + 1) : 2 * (S + 1) + S]
1085 return (Bs == f_Bs).long().min(dim=-1).values > 0
1087 def generate_prompts_and_answers_(self, nb, tasks=None, progress_bar=False):
1089 tasks = self.all_tasks
1091 S = self.height * self.width
1092 prompts = torch.zeros(nb, 3 * S + 2, dtype=torch.int64)
1093 answers = torch.zeros(nb, S, dtype=torch.int64)
1095 bunch = zip(prompts, answers)
1101 desc="world generation",
1102 total=prompts.size(0),
1105 for prompt, answer in bunch:
1106 A = prompt[0 * (S + 1) : 0 * (S + 1) + S].view(self.height, self.width)
1107 f_A = prompt[1 * (S + 1) : 1 * (S + 1) + S].view(self.height, self.width)
1108 B = prompt[2 * (S + 1) : 2 * (S + 1) + S].view(self.height, self.width)
1109 f_B = answer.view(self.height, self.width)
1110 task = tasks[torch.randint(len(tasks), (1,)).item()]
1111 task(A, f_A, B, f_B)
1113 return prompts.flatten(1), answers.flatten(1)
1115 def save_quiz_illustrations(
1121 predicted_prompts=None,
1122 predicted_answers=None,
1127 filename_prefix + ".png",
1135 def save_some_examples(self, result_dir):
1137 for t in self.all_tasks:
1139 prompts, answers = self.generate_prompts_and_answers_(nb, tasks=[t])
1140 self.save_quiz_illustrations(
1141 result_dir, t.__name__, prompts[:nb], answers[:nb], nrow=nrow
1145 ######################################################################
1147 if __name__ == "__main__":
1150 # grids = Grids(max_nb_cached_chunks=5, chunk_size=100, nb_threads=4)
1154 # grids = problem.MultiThreadProblem(
1155 # grids, max_nb_cached_chunks=50, chunk_size=100, nb_threads=1
1158 # start_time = time.perf_counter()
1159 # prompts, answers = grids.generate_prompts_and_answers(nb)
1160 # delay = time.perf_counter() - start_time
1161 # print(f"{prompts.size(0)/delay:02f} seq/s")
1168 # for t in grids.all_tasks:
1169 for t in [grids.task_islands]:
1171 prompts, answers = grids.generate_prompts_and_answers_(nb, tasks=[t])
1172 grids.save_quiz_illustrations(
1173 "/tmp", t.__name__, prompts[:nb], answers[:nb], nrow=nrow
1180 # for t in grids.all_tasks:
1181 for t in [grids.task_islands]:
1182 start_time = time.perf_counter()
1183 prompts, answers = grids.generate_prompts_and_answers_(nb, tasks=[t])
1184 delay = time.perf_counter() - start_time
1185 print(f"{t.__name__} {prompts.size(0)/delay:02f} seq/s")
1189 m = torch.randint(2, (prompts.size(0),))
1190 predicted_prompts = m * (torch.randint(2, (prompts.size(0),)) * 2 - 1)
1191 predicted_answers = (1 - m) * (torch.randint(2, (prompts.size(0),)) * 2 - 1)
1193 grids.save_quiz_illustrations(
1198 # You can add a bool to put a frame around the predicted parts
1199 predicted_prompts[:nb],
1200 predicted_answers[:nb],