("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]),
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
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,
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)
return y
- margin = 8
-
img_prompts = torch.cat(
[
add_frame(
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
)
]
+ # @torch.compile
def rec_coo_(self, x, n, min_height=3, min_width=3):
collision = x.new(x.size())
while True:
######################################################################
+ # @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
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
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
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
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
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
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
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)]:
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)]:
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
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
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):
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()
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)
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)
projects / culture.git / blobdiff