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_move(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     ######################################################################
 410
 411     def generate_prompts_and_answers(self, nb, device="cpu"):
 412         tasks = [
 413             self.task_replace_color,
 414             self.task_move,
 415             self.task_grow,
 416             self.task_color_grow,
 417             self.task_frame,
 418             self.task_detect,
 419         ]
 420         prompts = torch.zeros(nb, self.height, self.width * 3, dtype=torch.int64)
 421         answers = torch.zeros(nb, self.height, self.width, dtype=torch.int64)
 422         w = self.width
 423
 424         for prompt, answer in tqdm.tqdm(
 425             zip(prompts, answers),
 426             dynamic_ncols=True,
 427             desc="world generation",
 428             total=prompts.size(0),
 429         ):
 430             A = prompt[:, 0 * w : 1 * w]
 431             f_A = prompt[:, 1 * w : 2 * w]
 432             B = prompt[:, 2 * w : 3 * w]
 433             f_B = answer
 434             task = tasks[torch.randint(len(tasks), (1,))]
 435             task(A, f_A, B, f_B)
 436
 437         return prompts.flatten(1), answers.flatten(1)
 438
 439     def save_quizzes(
 440         self,
 441         result_dir,
 442         filename_prefix,
 443         prompts,
 444         answers,
 445         predicted_prompts=None,
 446         predicted_answers=None,
 447     ):
 448         self.save_image(
 449             result_dir,
 450             filename_prefix + ".png",
 451             prompts,
 452             answers,
 453             predicted_prompts,
 454             predicted_answers,
 455         )
 456
 457
 458 ######################################################################
 459
 460 if __name__ == "__main__":
 461     import time
 462
 463     reasoning = Reasoning()
 464
 465     start_time = time.perf_counter()
 466     prompts, answers = reasoning.generate_prompts_and_answers(100)
 467     delay = time.perf_counter() - start_time
 468     print(f"{prompts.size(0)/delay:02f} seq/s")
 469
 470     predicted_prompts = torch.rand(prompts.size(0)) < 0.5
 471     predicted_answers = torch.logical_not(predicted_prompts)
 472
 473     reasoning.save_quizzes(
 474         "/tmp",
 475         "test",
 476         prompts[:64],
 477         answers[:64],
 478         # You can add a bool to put a frame around the predicted parts
 479         predicted_prompts[:64],
 480         predicted_answers[:64],
 481     )