reasoning.py

   1 #!/usr/bin/env python
   2
   3 # Any copyright is dedicated to the Public Domain.
   4 # https://creativecommons.org/publicdomain/zero/1.0/
   5
   6 # Written by Francois Fleuret <francois@fleuret.org>
   7
   8 import math, sys, tqdm, os, warnings
   9
  10 import torch, torchvision
  11
  12 from torch import nn
  13 from torch.nn import functional as F
  14
  15 ######################################################################
  16
  17 import problem
  18
  19
  20 class Reasoning(problem.Problem):
  21     named_colors = [
  22         ("white", [255, 255, 255]),
  23         ("red", [255, 0, 0]),
  24         ("green", [0, 192, 0]),
  25         ("blue", [0, 0, 255]),
  26         ("orange", [255, 192, 0]),
  27         ("cyan", [0, 255, 255]),
  28         ("violet", [255, 0, 255]),
  29         ("lightgreen", [192, 255, 192]),
  30         ("brown", [165, 42, 42]),
  31         ("lightblue", [192, 192, 255]),
  32         ("gray", [128, 128, 128]),
  33     ]
  34
  35     def __init__(self, device=torch.device("cpu")):
  36         self.colors = torch.tensor([c for _, c in self.named_colors])
  37         self.name2color = dict([(p[0], i) for i, p in enumerate(self.named_colors)])
  38         self.height = 10
  39         self.width = 10
  40         self.device = device
  41
  42     ######################################################################
  43
  44     def frame2img(self, x, scale=15):
  45         x = x.reshape(x.size(0), self.height, -1)
  46         m = torch.logical_and(x >= 0, x < self.nb_token_values()).long()
  47         x = self.colors[x * m].permute(0, 3, 1, 2)
  48         s = x.shape
  49         x = x[:, :, :, None, :, None].expand(-1, -1, -1, scale, -1, scale)
  50         x = x.reshape(s[0], s[1], s[2] * scale, s[3] * scale)
  51
  52         x[:, :, :, torch.arange(0, x.size(3), scale)] = 0
  53         x[:, :, torch.arange(0, x.size(2), scale), :] = 0
  54         x = x[:, :, 1:, 1:]
  55
  56         for n in range(m.size(0)):
  57             for i in range(m.size(1)):
  58                 for j in range(m.size(2)):
  59                     if m[n, i, j] == 0:
  60                         for k in range(2, scale - 2):
  61                             for l in [0, 1]:
  62                                 x[n, :, i * scale + k, j * scale + k - l] = 0
  63                                 x[
  64                                     n, :, i * scale + scale - 1 - k, j * scale + k - l
  65                                 ] = 0
  66
  67         return x
  68
  69     def frame2img_(self, x, scale=15):
  70         x = x.reshape(x.size(0), self.height, -1)
  71         x = self.colors[x].permute(0, 3, 1, 2)
  72         s = x.shape
  73         x = x[:, :, :, None, :, None].expand(-1, -1, -1, scale, -1, scale)
  74         x = x.reshape(s[0], s[1], s[2] * scale, s[3] * scale)
  75
  76         x[:, :, :, torch.arange(0, x.size(3), scale)] = 0
  77         x[:, :, torch.arange(0, x.size(2), scale), :] = 0
  78         x = x[:, :, 1:, 1:]
  79
  80         return x
  81
  82     def save_image(
  83         self,
  84         result_dir,
  85         filename,
  86         prompts,
  87         answers,
  88         predicted_prompts=None,
  89         predicted_answers=None,
  90     ):
  91         prompts = prompts.reshape(prompts.size(0), self.height, -1)
  92         answers = answers.reshape(answers.size(0), self.height, -1)
  93
  94         if predicted_prompts is None:
  95             predicted_prompts = 255
  96
  97         if predicted_answers is None:
  98             predicted_answers = 255
  99
 100         def add_frame(x, c, margin, bottom=False):
 101             if bottom:
 102                 h, w, di, dj = x.size(2) + margin, x.size(3), 0, 0
 103             else:
 104                 h, w, di, dj = (
 105                     x.size(2) + 2 * margin,
 106                     x.size(3) + 2 * margin,
 107                     margin,
 108                     margin,
 109                 )
 110
 111             y = x.new_full((x.size(0), x.size(1), h, w), 0)
 112
 113             if type(c) is int:
 114                 y[...] = c
 115             else:
 116                 c = c.long()[:, None]
 117                 c = c * torch.tensor([192, 192, 192], device=c.device) + (
 118                     1 - c
 119                 ) * torch.tensor([255, 255, 255], device=c.device)
 120                 y[...] = c[:, :, None, None]
 121
 122             y[:, :, di : di + x.size(2), dj : dj + x.size(3)] = x
 123
 124             return y
 125
 126         margin = 8
 127
 128         img_prompts = torch.cat(
 129             [
 130                 add_frame(
 131                     add_frame(self.frame2img(x), c=0, margin=1),
 132                     c=predicted_prompts,
 133                     margin=margin,
 134                 )
 135                 for x in prompts.to("cpu").split(split_size=self.width, dim=2)
 136             ],
 137             dim=3,
 138         )
 139
 140         h = img_prompts.size(2)
 141         img_answers = add_frame(
 142             add_frame(self.frame2img(answers.to("cpu")), c=0, margin=1),
 143             c=predicted_answers,
 144             margin=margin,
 145         )
 146
 147         separator_size = 2 * margin
 148
 149         separator = img_prompts.new_full(
 150             (
 151                 img_prompts.size(0),
 152                 img_prompts.size(1),
 153                 img_prompts.size(2),
 154                 separator_size,
 155             ),
 156             255,
 157         )
 158
 159         marker = img_prompts.new_full(
 160             (
 161                 img_prompts.size(0),
 162                 img_prompts.size(1),
 163                 img_prompts.size(2),
 164                 separator_size,
 165             ),
 166             255,
 167         )
 168
 169         # marker[:, :, 0] = 0
 170         # marker[:, :, h - 1] = 0
 171
 172         for k in range(1, 2 * separator_size - 8):
 173             i = k - (separator_size - 4)
 174             j = separator_size - 5 - abs(i)
 175             marker[:, :, h // 2 - 1 + i, 2 + j] = 0
 176             marker[:, :, h // 2 - 1 + i + 1, 2 + j] = 0
 177
 178         img = torch.cat(
 179             [
 180                 img_prompts,
 181                 marker,
 182                 img_answers,
 183             ],
 184             dim=3,
 185         )
 186
 187         image_name = os.path.join(result_dir, filename)
 188         torchvision.utils.save_image(
 189             img.float() / 255.0, image_name, nrow=4, padding=margin * 4, pad_value=1.0
 190         )
 191
 192     ######################################################################
 193
 194     def nb_token_values(self):
 195         return len(self.colors)
 196
 197     # That's quite a tensorial spaghetti mess to sample
 198     # non-overlapping rectangles quickly, but made the generation of
 199     # 100k samples go from 1h50 with a lame pure python code to 3min30s
 200     # with this one.
 201     def rec_coo(self, nb_rec, min_height=3, min_width=3):
 202         nb_trials = 200
 203
 204         while True:
 205             v = (
 206                 (
 207                     torch.rand(nb_trials * nb_rec, self.height + 1, device=self.device)
 208                     .sort(dim=-1)
 209                     .indices
 210                     < 2
 211                 )
 212                 .long()
 213                 .cumsum(dim=1)
 214                 == 1
 215             ).long()
 216
 217             h = (
 218                 (
 219                     torch.rand(nb_trials * nb_rec, self.width + 1, device=self.device)
 220                     .sort(dim=-1)
 221                     .indices
 222                     < 2
 223                 )
 224                 .long()
 225                 .cumsum(dim=1)
 226                 == 1
 227             ).long()
 228
 229             i = torch.logical_and(
 230                 v.sum(dim=-1) >= min_height, h.sum(dim=-1) >= min_width
 231             )
 232
 233             v, h = v[i], h[i]
 234             v = v[: v.size(0) - v.size(0) % nb_rec]
 235             h = h[: h.size(0) - h.size(0) % nb_rec]
 236             v = v.reshape(v.size(0) // nb_rec, nb_rec, -1)
 237             h = h.reshape(h.size(0) // nb_rec, nb_rec, -1)
 238
 239             r = v[:, :, :, None] * h[:, :, None, :]
 240
 241             valid = r.sum(dim=1).flatten(1).max(dim=-1).values == 1
 242
 243             v = v[valid]
 244             h = h[valid]
 245
 246             if v.size(0) > 0:
 247                 break
 248
 249         av = torch.arange(v.size(2), device=self.device)[None, :]
 250         ah = torch.arange(h.size(2), device=self.device)[None, :]
 251
 252         return [
 253             (i1.item(), j1.item(), i2.item() + 1, j2.item() + 1)
 254             for i1, j1, i2, j2 in zip(
 255                 v.size(2) - (v[0] * (v.size(2) - av)).max(dim=-1).values,
 256                 h.size(2) - (h[0] * (h.size(2) - ah)).max(dim=-1).values,
 257                 (v[0] * av).max(dim=-1).values,
 258                 (h[0] * ah).max(dim=-1).values,
 259             )
 260         ]
 261
 262     def rec_coo_(self, x, n, min_height=3, min_width=3):
 263         collision = x.new(x.size())
 264         while True:
 265             collision[...] = 0
 266             result = []
 267             for _ in range(n):
 268                 while True:
 269                     i1, i2 = torch.randint(x.size(0), (2,))
 270                     if i1 + min_height <= i2:
 271                         break
 272                 while True:
 273                     j1, j2 = torch.randint(x.size(1), (2,))
 274                     if j1 + min_width <= j2:
 275                         break
 276                 collision[i1:i2, j1:j2] += 1
 277                 if collision.max() > 1:
 278                     break
 279                 result.append((i1, j1, i2, j2))
 280             if collision.max() == 1:
 281                 break
 282         return result
 283
 284     ######################################################################
 285
 286     def task_replace_color(self, A, f_A, B, f_B):
 287         nb_rec = 3
 288         c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
 289         for X, f_X in [(A, f_A), (B, f_B)]:
 290             r = self.rec_coo(nb_rec)
 291             for n in range(nb_rec):
 292                 i1, j1, i2, j2 = r[n]
 293                 X[i1:i2, j1:j2] = c[n]
 294                 f_X[i1:i2, j1:j2] = c[n if n > 0 else -1]
 295
 296     def task_translate(self, A, f_A, B, f_B):
 297         di, dj = torch.randint(3, (2,)) - 1
 298         nb_rec = 3
 299         c = torch.randperm(len(self.colors) - 1)[:nb_rec] + 1
 300         for X, f_X in [(A, f_A), (B, f_B)]:
 301             while True:
 302                 r = self.rec_coo(nb_rec)
 303                 i1, j1, i2, j2 = r[nb_rec - 1]
 304                 if (
 305                     i1 + di >= 0
 306                     and i2 + di < X.size(0)
 307                     and j1 + dj >= 0
 308                     and j2 + dj < X.size(1)
 309                 ):
 310                     break
 311
 312             for n in range(nb_rec):
 313                 i1, j1, i2, j2 = r[n]
 314                 X[i1:i2, j1:j2] = c[n]
 315                 if n == nb_rec - 1:
 316                     f_X[i1 + di : i2 + di, j1 + dj : j2 + dj] = c[n]
 317                 else:
 318                     f_X[i1:i2, j1:j2] = c[n]
 319
 320     def task_grow(self, A, f_A, B, f_B):
 321         di, dj = torch.randint(2, (2,)) * 2 - 1
 322         nb_rec = 3
 323         c = torch.randperm(len(self.colors) - 1)[:nb_rec] + 1
 324         direction = torch.randint(2, (1,))
 325         for X, f_X in [(A, f_A), (B, f_B)]:
 326             while True:
 327                 r = self.rec_coo(nb_rec)
 328                 i1, j1, i2, j2 = r[nb_rec - 1]
 329                 if i1 + 3 < i2 and j1 + 3 < j2:
 330                     break
 331
 332             for n in range(nb_rec):
 333                 i1, j1, i2, j2 = r[n]
 334                 if n == nb_rec - 1:
 335                     if direction == 0:
 336                         X[i1 + 1 : i2 - 1, j1 + 1 : j2 - 1] = c[n]
 337                         f_X[i1:i2, j1:j2] = c[n]
 338                     else:
 339                         X[i1:i2, j1:j2] = c[n]
 340                         f_X[i1 + 1 : i2 - 1, j1 + 1 : j2 - 1] = c[n]
 341                 else:
 342                     X[i1:i2, j1:j2] = c[n]
 343                     f_X[i1:i2, j1:j2] = c[n]
 344
 345     def task_color_grow(self, A, f_A, B, f_B):
 346         di, dj = torch.randint(2, (2,)) * 2 - 1
 347         nb_rec = 3
 348         c = torch.randperm(len(self.colors) - 1)[: 2 * nb_rec] + 1
 349         direction = torch.randint(4, (1,))
 350         for X, f_X in [(A, f_A), (B, f_B)]:
 351             r = self.rec_coo(nb_rec)
 352             for n in range(nb_rec):
 353                 i1, j1, i2, j2 = r[n]
 354                 X[i1:i2, j1:j2] = c[2 * n]
 355                 f_X[i1:i2, j1:j2] = c[2 * n]
 356                 # Not my proudest moment
 357                 if direction == 0:
 358                     i = (i1 + i2) // 2
 359                     X[i : i + 1, j1:j2] = c[2 * n + 1]
 360                     if n == nb_rec - 1:
 361                         f_X[i:i2, j1:j2] = c[2 * n + 1]
 362                     else:
 363                         f_X[i : i + 1, j1:j2] = c[2 * n + 1]
 364                 elif direction == 1:
 365                     i = (i1 + i2 - 1) // 2
 366                     X[i : i + 1, j1:j2] = c[2 * n + 1]
 367                     if n == nb_rec - 1:
 368                         f_X[i1 : i + 1, j1:j2] = c[2 * n + 1]
 369                     else:
 370                         f_X[i : i + 1, j1:j2] = c[2 * n + 1]
 371                 elif direction == 2:
 372                     j = (j1 + j2) // 2
 373                     X[i1:i2, j : j + 1] = c[2 * n + 1]
 374                     if n == nb_rec - 1:
 375                         f_X[i1:i2, j:j2] = c[2 * n + 1]
 376                     else:
 377                         f_X[i1:i2, j : j + 1] = c[2 * n + 1]
 378                 elif direction == 3:
 379                     j = (j1 + j2 - 1) // 2
 380                     X[i1:i2, j : j + 1] = c[2 * n + 1]
 381                     if n == nb_rec - 1:
 382                         f_X[i1:i2, j1 : j + 1] = c[2 * n + 1]
 383                     else:
 384                         f_X[i1:i2, j : j + 1] = c[2 * n + 1]
 385
 386     def task_frame(self, A, f_A, B, f_B):
 387         nb_rec = 3
 388         c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
 389         for X, f_X in [(A, f_A), (B, f_B)]:
 390             r = self.rec_coo(nb_rec)
 391             for n in range(nb_rec):
 392                 i1, j1, i2, j2 = r[n]
 393                 X[i1:i2, j1:j2] = c[n]
 394                 f_X[i1:i2, j1:j2] = c[n]
 395                 if n == nb_rec - 1:
 396                     f_X[i1 + 1 : i2 - 1, j1 + 1 : j2 - 1] = 0
 397
 398     def task_detect(self, A, f_A, B, f_B):
 399         nb_rec = 3
 400         c = torch.randperm(len(self.colors) - 1)[: nb_rec + 1] + 1
 401         for X, f_X in [(A, f_A), (B, f_B)]:
 402             r = self.rec_coo(nb_rec)
 403             for n in range(nb_rec):
 404                 i1, j1, i2, j2 = r[n]
 405                 X[i1:i2, j1:j2] = c[n]
 406                 if n < nb_rec - 1:
 407                     f_X[i1, j1] = c[-1]
 408
 409     def task_count(self, A, f_A, B, f_B):
 410         N = torch.randint(4, (1,)) + 2
 411         c = torch.randperm(len(self.colors) - 1)[:N] + 1
 412
 413         for X, f_X in [(A, f_A), (B, f_B)]:
 414
 415             def contact(i, j, q):
 416                 nq, nq_diag = 0, 0
 417                 no = 0
 418
 419                 for ii, jj in [
 420                     (i - 1, j - 1),
 421                     (i - 1, j),
 422                     (i - 1, j + 1),
 423                     (i, j - 1),
 424                     (i, j + 1),
 425                     (i + 1, j - 1),
 426                     (i + 1, j),
 427                     (i + 1, j + 1),
 428                 ]:
 429                     if ii >= 0 and ii < self.height and jj >= 0 and jj < self.width:
 430                         if X[ii, jj] != 0 and X[ii, jj] != q:
 431                             no += 1
 432
 433                 for ii, jj in [
 434                     (i - 1, j - 1),
 435                     (i - 1, j + 1),
 436                     (i + 1, j - 1),
 437                     (i + 1, j + 1),
 438                 ]:
 439                     if ii >= 0 and ii < self.height and jj >= 0 and jj < self.width:
 440                         if X[ii, jj] == q and X[i, jj] != q and X[ii, j] != q:
 441                             nq_diag += 1
 442
 443                 for ii, jj in [(i - 1, j), (i, j - 1), (i, j + 1), (i + 1, j)]:
 444                     if ii >= 0 and ii < self.height and jj >= 0 and jj < self.width:
 445                         if X[ii, jj] == q:
 446                             nq += 1
 447
 448                 return no, nq, nq_diag
 449
 450             nb = torch.zeros(N, dtype=torch.int64)
 451             q = torch.randint(N, (self.height * self.width,))
 452             k = torch.randperm(self.height * self.width)
 453             for p in range(self.height * self.width):
 454                 i, j = k[p] % self.height, k[p] // self.height
 455                 no, nq, nq_diag = contact(i, j, c[q[p]])
 456                 if no == 0 and nq_diag == 0:
 457                     if nq == 0:
 458                         if nb[q[p]] < self.width:
 459                             X[i, j] = c[q[p]]
 460                             nb[q[p]] += 1
 461                     if nq == 1:
 462                         X[i, j] = c[q[p]]
 463
 464             for n in range(N):
 465                 for j in range(nb[n]):
 466                     f_X[n, j] = c[n]
 467
 468     def task_trajectory(self, A, f_A, B, f_B):
 469         c = torch.randperm(len(self.colors) - 1)[:2] + 1
 470         for X, f_X in [(A, f_A), (B, f_B)]:
 471             while True:
 472                 di, dj = torch.randint(7, (2,)) - 3
 473                 i, j = torch.randint(self.height, (1,)), torch.randint(self.width, (1,))
 474                 if (
 475                     abs(di) + abs(dj) > 0
 476                     and i + 2 * di >= 0
 477                     and i + 2 * di < self.height
 478                     and j + 2 * dj >= 0
 479                     and j + 2 * dj < self.width
 480                 ):
 481                     break
 482
 483             k = 0
 484             while (
 485                 i + k * di >= 0
 486                 and i + k * di < self.height
 487                 and j + k * dj >= 0
 488                 and j + k * dj < self.width
 489             ):
 490                 if k < 2:
 491                     X[i + k * di, j + k * dj] = c[k]
 492                 f_X[i + k * di, j + k * dj] = c[min(k, 1)]
 493                 k += 1
 494
 495     def task_bounce(self, A, f_A, B, f_B):
 496         c = torch.randperm(len(self.colors) - 1)[:3] + 1
 497         for X, f_X in [(A, f_A), (B, f_B)]:
 498
 499             def free(i, j):
 500                 return (
 501                     i >= 0
 502                     and i < self.height
 503                     and j >= 0
 504                     and j < self.width
 505                     and f_X[i, j] == 0
 506                 )
 507
 508             while True:
 509                 f_X[...] = 0
 510                 X[...] = 0
 511
 512                 for _ in range((self.height * self.width) // 10):
 513                     i, j = torch.randint(self.height, (1,)), torch.randint(
 514                         self.width, (1,)
 515                     )
 516                     X[i, j] = c[0]
 517                     f_X[i, j] = c[0]
 518
 519                 while True:
 520                     di, dj = torch.randint(7, (2,)) - 3
 521                     if abs(di) + abs(dj) == 1:
 522                         break
 523
 524                 i, j = torch.randint(self.height, (1,)), torch.randint(self.width, (1,))
 525
 526                 X[i, j] = c[1]
 527                 f_X[i, j] = c[1]
 528                 l = 0
 529
 530                 while True:
 531                     l += 1
 532                     if free(i + di, j + dj):
 533                         pass
 534                     elif free(i - dj, j + di):
 535                         di, dj = -dj, di
 536                         if free(i + dj, j - di):
 537                             if torch.rand(1) < 0.5:
 538                                 di, dj = -di, -dj
 539                     elif free(i + dj, j - di):
 540                         di, dj = dj, -di
 541                     else:
 542                         break
 543
 544                     i, j = i + di, j + dj
 545                     f_X[i, j] = c[2]
 546                     if l <= 1:
 547                         X[i, j] = c[2]
 548
 549                     if l >= self.width:
 550                         break
 551
 552                 f_X[i, j] = c[1]
 553                 X[i, j] = c[1]
 554
 555                 if l > 3:
 556                     break
 557
 558     ######################################################################
 559
 560     def generate_prompts_and_answers(self, nb, device="cpu"):
 561         tasks = [
 562             self.task_replace_color,
 563             self.task_translate,
 564             self.task_grow,
 565             self.task_color_grow,
 566             self.task_frame,
 567             self.task_detect,
 568             self.task_count,
 569             self.task_trajectory,
 570             self.task_bounce,
 571         ]
 572         prompts = torch.zeros(nb, self.height, self.width * 3, dtype=torch.int64)
 573         answers = torch.zeros(nb, self.height, self.width, dtype=torch.int64)
 574         w = self.width
 575
 576         for prompt, answer in tqdm.tqdm(
 577             zip(prompts, answers),
 578             dynamic_ncols=True,
 579             desc="world generation",
 580             total=prompts.size(0),
 581         ):
 582             A = prompt[:, 0 * w : 1 * w]
 583             f_A = prompt[:, 1 * w : 2 * w]
 584             B = prompt[:, 2 * w : 3 * w]
 585             f_B = answer
 586             task = tasks[torch.randint(len(tasks), (1,))]
 587             task(A, f_A, B, f_B)
 588
 589         return prompts.flatten(1), answers.flatten(1)
 590
 591     def save_quizzes(
 592         self,
 593         result_dir,
 594         filename_prefix,
 595         prompts,
 596         answers,
 597         predicted_prompts=None,
 598         predicted_answers=None,
 599     ):
 600         self.save_image(
 601             result_dir,
 602             filename_prefix + ".png",
 603             prompts,
 604             answers,
 605             predicted_prompts,
 606             predicted_answers,
 607         )
 608
 609
 610 ######################################################################
 611
 612 if __name__ == "__main__":
 613     import time
 614
 615     reasoning = Reasoning()
 616
 617     start_time = time.perf_counter()
 618     prompts, answers = reasoning.generate_prompts_and_answers(100)
 619     delay = time.perf_counter() - start_time
 620     print(f"{prompts.size(0)/delay:02f} seq/s")
 621
 622     predicted_prompts = torch.rand(prompts.size(0)) < 0.5
 623     predicted_answers = torch.logical_not(predicted_prompts)
 624
 625     reasoning.save_quizzes(
 626         "/tmp",
 627         "test",
 628         prompts[:64],
 629         answers[:64],
 630         # You can add a bool to put a frame around the predicted parts
 631         # predicted_prompts[:64],
 632         # predicted_answers[:64],
 633     )