max_nb_cached_chunks=None,
chunk_size=None,
nb_threads=-1,
+ tasks=None,
):
self.colors = torch.tensor([c for _, c in self.named_colors])
self.height = 10
self.width = 10
self.cache_rec_coo = {}
+
+ all_tasks = [
+ self.task_replace_color,
+ self.task_translate,
+ self.task_grow,
+ self.task_half_fill,
+ self.task_frame,
+ self.task_detect,
+ self.task_count,
+ self.task_trajectory,
+ self.task_bounce,
+ self.task_scale,
+ self.task_symbols,
+ self.task_isometry,
+ # self.task_path,
+ ]
+
+ if tasks is None:
+ self.all_tasks = all_tasks
+ else:
+ self.all_tasks = [getattr(self, "task_" + t) for t in tasks.split(",")]
+
super().__init__(max_nb_cached_chunks, chunk_size, nb_threads)
######################################################################
f_X[i1:i2, j1:j2] = c[n]
# @torch.compile
- def task_color_grow(self, A, f_A, B, f_B):
+ def task_half_fill(self, A, f_A, B, f_B):
di, dj = torch.randint(2, (2,)) * 2 - 1
nb_rec = 3
c = torch.randperm(len(self.colors) - 1)[: 2 * nb_rec] + 1
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):
+ def task_isometry(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]])
# i,j=q%self.height,q//self.height
# if
- ######################################################################
+ # @torch.compile
+ def task_puzzle(self, A, f_A, B, f_B):
+ S = 4
+ i0, j0 = (self.height - S) // 2, (self.width - S) // 2
+ c = torch.randperm(len(self.colors) - 1)[:4] + 1
+ for X, f_X in [(A, f_A), (B, f_B)]:
+ while True:
+ f_X[...] = 0
+ h = list(torch.randperm(c.size(0)))
+ n = torch.zeros(c.max() + 1)
+ for _ in range(2):
+ k = torch.randperm(S * S)
+ for q in k:
+ i, j = q % S + i0, q // S + j0
+ if f_X[i, j] == 0:
+ r, s, t, u = (
+ f_X[i - 1, j],
+ f_X[i, j - 1],
+ f_X[i + 1, j],
+ f_X[i, j + 1],
+ )
+ r, s, t, u = torch.tensor([r, s, t, u])[torch.randperm(4)]
+ if r > 0 and n[r] < 6:
+ n[r] += 1
+ f_X[i, j] = r
+ elif s > 0 and n[s] < 6:
+ n[s] += 1
+ f_X[i, j] = s
+ elif t > 0 and n[t] < 6:
+ n[t] += 1
+ f_X[i, j] = t
+ elif u > 0 and n[u] < 6:
+ n[u] += 1
+ f_X[i, j] = u
+ else:
+ if len(h) > 0:
+ d = c[h.pop()]
+ n[d] += 1
+ f_X[i, j] = d
+
+ if n.sum() == S * S:
+ break
- def all_tasks(self):
- return [
- self.task_replace_color,
- self.task_translate,
- self.task_grow,
- self.task_color_grow,
- self.task_frame,
- self.task_detect,
- self.task_count,
- self.task_trajectory,
- self.task_bounce,
- self.task_scale,
- self.task_symbols,
- self.task_ortho,
- # self.task_path,
- ]
+ k = 0
+ for d in range(4):
+ while True:
+ ii, jj = torch.randint(self.height, (1,)), torch.randint(
+ self.width, (1,)
+ )
+ e = 0
+ for i in range(S):
+ for j in range(S):
+ if (
+ ii + i >= self.height
+ or jj + j >= self.width
+ or (
+ f_X[i + i0, j + j0] == c[d]
+ and X[ii + i, jj + j] > 0
+ )
+ ):
+ e = 1
+ if e == 0:
+ break
+ for i in range(S):
+ for j in range(S):
+ if f_X[i + i0, j + j0] == c[d]:
+ X[ii + i, jj + j] = c[d]
+
+ ######################################################################
def trivial_prompts_and_answers(self, prompts, answers):
S = self.height * self.width
def generate_prompts_and_answers_(self, nb, tasks=None, progress_bar=False):
if tasks is None:
- tasks = self.all_tasks()
+ tasks = self.all_tasks
S = self.height * self.width
prompts = torch.zeros(nb, 3 * S + 2, dtype=torch.int64)
def save_some_examples(self, result_dir):
nb, nrow = 72, 4
- for t in self.all_tasks():
+ for t in self.all_tasks:
print(t.__name__)
prompts, answers = self.generate_prompts_and_answers_(nb, tasks=[t])
self.save_quizzes(
nb, nrow = 72, 4
# nb, nrow = 8, 2
- # for t in grids.all_tasks():
- for t in [grids.task_path]:
+ # for t in grids.all_tasks:
+ for t in [
+ grids.task_replace_color,
+ grids.task_frame,
+ ]:
print(t.__name__)
prompts, answers = grids.generate_prompts_and_answers_(nb, tasks=[t])
grids.save_quizzes("/tmp", t.__name__, prompts[:nb], answers[:nb], nrow=nrow)
- # exit(0)
+ exit(0)
nb = 1000
- for t in grids.all_tasks():
- # for t in [ grids.task_replace_color ]: #grids.all_tasks():
+ for t in grids.all_tasks:
+ # for t in [ grids.task_replace_color ]: #grids.all_tasks:
start_time = time.perf_counter()
prompts, answers = grids.generate_prompts_and_answers_(nb, tasks=[t])
delay = time.perf_counter() - start_time