main.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 # torch.backends.cuda.matmul.allow_tf23
   9 # torch.autocast(torch.bfloat16)
  10
  11 import math, sys, argparse, time, tqdm, os
  12
  13 import torch, torchvision
  14 from torch import nn
  15 from torch.nn import functional as F
  16
  17 import mygpt, tensorstack
  18
  19 ######################################################################
  20
  21 if torch.cuda.is_available():
  22     device = torch.device("cuda")
  23     torch.backends.cuda.matmul.allow_tf32 = True
  24 else:
  25     device = torch.device("cpu")
  26
  27 ######################################################################
  28
  29 parser = argparse.ArgumentParser(
  30     description="An implementation of GPT with cache.",
  31     formatter_class=argparse.ArgumentDefaultsHelpFormatter,
  32 )
  33
  34 parser.add_argument(
  35     "--task", type=str, default="picoclvr", help="picoclvr, mnist, maze, snake, stack, expr"
  36 )
  37
  38 parser.add_argument("--log_filename", type=str, default="train.log", help=" ")
  39
  40 parser.add_argument("--result_dir", type=str, default=None)
  41
  42 parser.add_argument("--seed", type=int, default=0)
  43
  44 parser.add_argument("--nb_epochs", type=int, default=None)
  45
  46 parser.add_argument("--batch_size", type=int, default=None)
  47
  48 parser.add_argument("--nb_train_samples", type=int, default=None)
  49
  50 parser.add_argument("--nb_test_samples", type=int, default=None)
  51
  52 parser.add_argument("--optim", type=str, default="adam")
  53
  54 parser.add_argument("--learning_rate", type=float, default=1e-4)
  55
  56 parser.add_argument("--learning_rate_schedule", type=str, default="10: 2e-5,30: 4e-6")
  57
  58 parser.add_argument("--dim_model", type=int, default=512)
  59
  60 parser.add_argument("--dim_keys", type=int, default=64)
  61
  62 parser.add_argument("--dim_hidden", type=int, default=2048)
  63
  64 parser.add_argument("--nb_heads", type=int, default=8)
  65
  66 parser.add_argument("--nb_blocks", type=int, default=12)
  67
  68 parser.add_argument("--dropout", type=float, default=0.1)
  69
  70 parser.add_argument("--deterministic_synthesis", action="store_true", default=False)
  71
  72 parser.add_argument("--no_checkpoint", action="store_true", default=False)
  73
  74 parser.add_argument("--overwrite_results", action="store_true", default=False)
  75
  76 parser.add_argument("--checkpoint_name", type=str, default="checkpoint.pth")
  77
  78 ##############################
  79 # picoclvr options
  80
  81 parser.add_argument("--picoclvr_nb_colors", type=int, default=5)
  82
  83 parser.add_argument("--picoclvr_height", type=int, default=12)
  84
  85 parser.add_argument("--picoclvr_width", type=int, default=16)
  86
  87 parser.add_argument("--picocvlr_prune_properties", type=str, default="none")
  88
  89 ##############################
  90 # Maze options
  91
  92 parser.add_argument("--maze_height", type=int, default=13)
  93
  94 parser.add_argument("--maze_width", type=int, default=21)
  95
  96 parser.add_argument("--maze_nb_walls", type=int, default=15)
  97
  98 ##############################
  99 # Snake options
 100
 101 parser.add_argument("--snake_height", type=int, default=6)
 102
 103 parser.add_argument("--snake_width", type=int, default=8)
 104
 105 parser.add_argument("--snake_nb_colors", type=int, default=5)
 106
 107 parser.add_argument("--snake_length", type=int, default=200)
 108
 109 ##############################
 110 # Snake options
 111
 112 parser.add_argument("--stack_nb_steps", type=int, default=100)
 113
 114 parser.add_argument("--stack_nb_stacks", type=int, default=1)
 115
 116 parser.add_argument("--stack_nb_digits", type=int, default=3)
 117
 118 parser.add_argument("--stack_fraction_values_for_train", type=float, default=None)
 119
 120 ######################################################################
 121
 122 args = parser.parse_args()
 123
 124 assert args.picocvlr_prune_properties in {"none", "train+eval", "eval"}
 125
 126 if args.result_dir is None:
 127     args.result_dir = f"results_{args.task}"
 128
 129 ######################################################################
 130
 131 default_args = {
 132     "picoclvr": {
 133         "nb_epochs": 25,
 134         "batch_size": 25,
 135         "nb_train_samples": 250000,
 136         "nb_test_samples": 10000,
 137     },
 138     "mnist": {
 139         "nb_epochs": 25,
 140         "batch_size": 10,
 141         "nb_train_samples": 250000,
 142         "nb_test_samples": 10000,
 143     },
 144     "maze": {
 145         "nb_epochs": 25,
 146         "batch_size": 25,
 147         "nb_train_samples": 250000,
 148         "nb_test_samples": 10000,
 149     },
 150     "snake": {
 151         "nb_epochs": 5,
 152         "batch_size": 25,
 153         "nb_train_samples": 250000,
 154         "nb_test_samples": 10000,
 155     },
 156     "stack": {
 157         "nb_epochs": 5,
 158         "batch_size": 25,
 159         "nb_train_samples": 100000,
 160         "nb_test_samples": 1000,
 161     },
 162     "expr": {
 163         "nb_epochs": 5,
 164         "batch_size": 25,
 165         "nb_train_samples": 100000,
 166         "nb_test_samples": 1000,
 167     },
 168 }
 169
 170 if args.task in default_args:
 171     for k, v in default_args[args.task].items():
 172         if getattr(args, k) is None:
 173             setattr(args, k, v)
 174
 175 ######################################################################
 176
 177 try:
 178     os.mkdir(args.result_dir)
 179 except FileExistsError:
 180     if not args.overwrite_results:
 181         print(f"result directory {args.result_dir} already exists")
 182         exit(1)
 183
 184 log_file = open(os.path.join(args.result_dir, args.log_filename), "a")
 185
 186 if args.seed >= 0:
 187     # torch.backends.cudnn.deterministic = True
 188     # torch.backends.cudnn.benchmark = False
 189     # torch.use_deterministic_algorithms(True)
 190     torch.manual_seed(args.seed)
 191     if torch.cuda.is_available():
 192         torch.cuda.manual_seed_all(args.seed)
 193
 194 ######################################################################
 195
 196
 197 def log_string(s):
 198     t = time.strftime("%Y%m%d-%H:%M:%S ", time.localtime())
 199
 200     if log_file is not None:
 201         log_file.write(t + s + "\n")
 202         log_file.flush()
 203
 204     print(t + s)
 205     sys.stdout.flush()
 206
 207
 208 for n in vars(args):
 209     log_string(f"args.{n} {getattr(args, n)}")
 210
 211 ######################################################################
 212
 213
 214 # ra_mask is boolean, with 1s on the values to generate
 215
 216
 217 def masked_inplace_autoregression(
 218     model,
 219     batch_size,
 220     input,
 221     ar_mask,
 222     forbidden_tokens=None,
 223     progress_bar_desc="autoregression",
 224     device=torch.device("cpu"),
 225 ):
 226
 227     batches = zip(input.split(batch_size), ar_mask.split(batch_size))
 228
 229     if progress_bar_desc is not None:
 230         batches = tqdm.tqdm(
 231             batches,
 232             dynamic_ncols=True,
 233             desc=progress_bar_desc,
 234             total=input.size(0) // batch_size,
 235         )
 236
 237     for input, ar_mask in batches:
 238         i = (ar_mask.sum(0) > 0).nonzero()
 239         if i.min() > 0:
 240             model(
 241                 mygpt.BracketedSequence(input, 0, i.min())
 242             )  # Needed to initialize the model's cache
 243         for s in range(i.min(), i.max() + 1):
 244             output = model(mygpt.BracketedSequence(input, s, 1)).x
 245             logits = output[:, s]
 246             if forbidden_tokens is not None:
 247                 logits = logits.masked_fill(forbidden_tokens, float("-inf"))
 248             if args.deterministic_synthesis:
 249                 t_next = logits.argmax(1)
 250             else:
 251                 dist = torch.distributions.categorical.Categorical(logits=logits)
 252                 t_next = dist.sample()
 253             input[:, s] = ar_mask[:, s] * t_next + (1 - ar_mask[:, s]) * input[:, s]
 254
 255
 256 ######################################################################
 257
 258
 259 class Task:
 260     def batches(self, split="train"):
 261         pass
 262
 263     def vocabulary_size(self):
 264         pass
 265
 266     def produce_results(self, n_epoch, model):
 267         pass
 268
 269
 270 ######################################################################
 271
 272 import picoclvr
 273
 274
 275 class TaskPicoCLVR(Task):
 276     # Make a tensor from a list of strings
 277     def tensorize(self, descr):
 278         token_descr = [s.strip().split(" ") for s in descr]
 279         l = max([len(s) for s in token_descr])
 280         token_descr = [s + ["<nul>"] * (l - len(s)) for s in token_descr]
 281         id_descr = [[self.token2id[u] for u in s] for s in token_descr]
 282         return torch.tensor(id_descr, device=self.device)
 283
 284     # Make a list of strings from a tensor
 285     def detensorize(self, x):
 286         return [" ".join([self.id2token[t.item()] for t in r]) for r in x]
 287
 288     # trim all the tensors in the tuple z to remove as much token from
 289     # left and right in the first tensor. If z is a tuple, all its
 290     # elements are trimed according to the triming for the first
 291     def trim(self, z, token="<nul>"):
 292         n = self.token2id[token]
 293         if type(z) == tuple:
 294             x = z[0]
 295             i = (1 - (F.pad(x, (1, 1), value=n) == n).min(0).values.long()).cumsum(0)
 296             a, b = (i == 0).nonzero().max(), (i == i.max()).nonzero().min()
 297             return tuple([t[:, a:b] for t in z])
 298         else:
 299             i = (1 - (F.pad(z, (1, 1), value=n) == n).min(0).values.long()).cumsum(0)
 300             a, b = (i == 0).nonzero().max(), (i == i.max()).nonzero().min()
 301             return z[:, a:b]
 302
 303     ######################
 304     # Not the cleanest part of the code
 305
 306     # Extract the last image of each sequence, from the last <img>
 307     # included, and set to <nul> all the tokens from the beginning of
 308     # that image to the end
 309     def excise_last_image(self, input):
 310         t_img, t_nul = self.token2id["<img>"], self.token2id["<nul>"]
 311         nb_img_tokens = self.height * self.width + 1
 312
 313         input = input.clone()
 314         t = (input == t_img).long()
 315         tail_masks = (t.cumsum(dim=1) == t.sum(dim=1, keepdim=True)).long()
 316         i = (t * tail_masks).nonzero(as_tuple=True)
 317         j = (
 318             i[0][:, None],
 319             i[1][:, None] + torch.arange(nb_img_tokens, device=input.device)[None, :],
 320         )
 321         images = self.trim(input[j])
 322         input[j] = t_nul
 323         loss_masks = 1 - tail_masks
 324         input, loss_masks = self.trim((input, loss_masks))
 325         return input, loss_masks, images
 326
 327     def add_true_image(self, input, images, loss_masks):
 328         t_nul = self.token2id["<nul>"]
 329         nb_img_tokens = self.height * self.width + 1
 330         input = F.pad(input, (0, nb_img_tokens), value=t_nul)
 331         loss_masks = F.pad(loss_masks, (0, nb_img_tokens), value=0)
 332         t = (input == t_nul).long()
 333         i = (t.cumsum(dim=1) == 1).nonzero(as_tuple=True)
 334         j = (
 335             i[0][:, None],
 336             i[1][:, None] + torch.arange(nb_img_tokens, device=input.device)[None, :],
 337         )
 338         input[j] = images
 339         loss_masks[j] = 1
 340         input, loss_masks = self.trim((input, loss_masks))
 341         return input, loss_masks
 342
 343     def add_generated_image(self, input, loss_masks, model):
 344         t_img, t_nul = self.token2id["<img>"], self.token2id["<nul>"]
 345         nb_img_tokens = self.height * self.width + 1
 346
 347         input = F.pad(input, (0, nb_img_tokens), value=t_nul)
 348         loss_masks = F.pad(loss_masks, (0, nb_img_tokens), value=0)
 349         t = (input == t_nul).long()
 350         i = (t.cumsum(dim=1) == 1).nonzero(as_tuple=True)
 351         input[i] = t_img
 352
 353         j = (
 354             i[0][:, None],
 355             i[1][:, None]
 356             + 1
 357             + torch.arange(nb_img_tokens - 1, device=input.device)[None, :],
 358         )
 359         ar_masks = input.new_zeros(input.size(), dtype=torch.int64)
 360         ar_masks[j] = 1
 361         forbidden_tokens = (
 362             torch.arange(self.vocabulary_size(), device=input.device) == t_nul
 363         )
 364         with torch.autograd.no_grad():
 365             t = model.training
 366             model.eval()
 367             masked_inplace_autoregression(
 368                 model,
 369                 self.batch_size,
 370                 input,
 371                 ar_masks,
 372                 forbidden_tokens,
 373                 progress_bar_desc=None,
 374                 device=self.device,
 375             )
 376             model.train(t)
 377
 378         input, loss_masks = self.trim((input, loss_masks))
 379
 380         return input, loss_masks
 381
 382     ######################
 383
 384     def __init__(
 385         self,
 386         nb_train_samples,
 387         nb_test_samples,
 388         batch_size,
 389         height,
 390         width,
 391         nb_colors=5,
 392         device=torch.device("cpu"),
 393         pruner_train=None,
 394         pruner_eval=None,
 395     ):
 396         def generate_descr(nb, cache_suffix, pruner):
 397             return picoclvr.generate(
 398                 nb,
 399                 height=self.height,
 400                 width=self.width,
 401                 nb_colors=nb_colors,
 402                 pruner=pruner,
 403             )
 404
 405         self.height = height
 406         self.width = width
 407         self.batch_size = batch_size
 408         self.device = device
 409         self.pruner_train = pruner_train
 410         self.pruner_eval = pruner_eval
 411
 412         param = {
 413             "nb_train_samples": nb_train_samples,
 414             "nb_test_samples": nb_test_samples,
 415             "height": height,
 416             "width": width,
 417             "nb_colors": nb_colors,
 418             "batch_size": batch_size,
 419             "rng_state": list(torch.get_rng_state()),
 420         }
 421
 422         log_string(
 423             f"generating {nb_train_samples+nb_test_samples} samples (can take some time)"
 424         )
 425         self.train_descr = generate_descr(
 426             nb_train_samples, "train", pruner=self.pruner_train
 427         )
 428         self.test_descr = generate_descr(nb_test_samples, "test", pruner=None)
 429
 430         # Build the tokenizer
 431         tokens = {"<nul>", "<img>"}
 432         for d in [self.train_descr, self.test_descr]:
 433             for s in d:
 434                 for t in s.strip().split(" "):
 435                     tokens.add(t)
 436         # make this set a sorted list to get the same tensors given
 437         # the same descr
 438         tokens = list(tokens)
 439         tokens.sort()
 440         self.token2id = dict([(t, n) for n, t in enumerate(tokens)])
 441         self.id2token = dict([(n, t) for n, t in enumerate(tokens)])
 442
 443         # Tokenize the train and test sets
 444         self.train_input = self.tensorize(self.train_descr)
 445         self.test_input = self.tensorize(self.test_descr)
 446
 447     def batches(self, split="train"):
 448         assert split in {"train", "test"}
 449         input = self.train_input if split == "train" else self.test_input
 450         for batch in tqdm.tqdm(
 451             input.split(self.batch_size), dynamic_ncols=True, desc=f"epoch-{split}"
 452         ):
 453             yield self.trim(batch)
 454
 455     def vocabulary_size(self):
 456         return len(self.token2id)
 457
 458     def compute_missing_properties(self, n_epoch, model, pruner=None):
 459         acc_nb_requested_properties = []
 460         acc_nb_missing_properties = []
 461         acc_nb_results = 0
 462
 463         for input in tqdm.tqdm(
 464             self.test_input.split(self.batch_size),
 465             dynamic_ncols=True,
 466             desc=f"test-properties",
 467         ):
 468             tape, loss_masks, _ = self.excise_last_image(input)
 469             tape, loss_masks = self.add_generated_image(tape, loss_masks, model)
 470             result_descr = self.detensorize(tape)
 471             np = picoclvr.nb_properties(
 472                 result_descr,
 473                 height=self.height,
 474                 width=self.width,
 475                 pruner=pruner,
 476             )
 477             nb_requested_properties, _, nb_missing_properties = zip(*np)
 478             acc_nb_requested_properties += nb_requested_properties
 479             acc_nb_missing_properties += nb_missing_properties
 480             acc_nb_results += len(result_descr)
 481
 482         nb_requested_properties = sum(acc_nb_requested_properties)
 483         nb_missing_properties = sum(acc_nb_missing_properties)
 484
 485         prefix = "" if pruner is None else "pruned_"
 486         log_string(f"nb_{prefix}samples {n_epoch} {acc_nb_results}")
 487         log_string(
 488             f"property_{prefix}nb {n_epoch} requested {sum(acc_nb_requested_properties)} missing {sum(acc_nb_missing_properties)}"
 489         )
 490         log_string(
 491             f"property_{prefix}miss {n_epoch} {100*nb_missing_properties/nb_requested_properties:.02f}%"
 492         )
 493
 494     ######################################################################
 495
 496     def produce_results(self, n_epoch, model):
 497         self.compute_missing_properties(n_epoch, model)
 498
 499         if self.pruner_eval is not None:
 500             self.compute_missing_properties(n_epoch, model, self.pruner_eval)
 501
 502         nb_tokens_to_generate = self.height * self.width + 3
 503         result_descr = []
 504         nb_per_primer = 8
 505         primer = []
 506
 507         for primer_descr in [
 508             "red above green <sep> green top <sep> blue right of red",
 509             "there is red <sep> there is yellow <sep> there is blue",
 510             "red below yellow <sep> yellow below green <sep> green below blue <sep> red right <sep> yellow left <sep> green right <sep> blue left",
 511             "green bottom <sep> yellow bottom <sep> green left of blue <sep> yellow right of blue <sep> blue top",
 512         ]:
 513             primer += [primer_descr] * nb_per_primer
 514
 515         tape = self.tensorize(primer)
 516         loss_masks = 1 - (tape == self.token2id["<nul>"]).long()
 517         tape, loss_masks = self.add_generated_image(tape, loss_masks, model)
 518         result_descr = self.detensorize(tape)
 519
 520         np = picoclvr.nb_properties(result_descr, height=self.height, width=self.width)
 521
 522         acc_nb_requested_properties, _, acc_nb_missing_properties = zip(*np)
 523         acc_nb_results = len(result_descr)
 524
 525         nb_requested_properties = sum(acc_nb_requested_properties)
 526         nb_missing_properties = sum(acc_nb_missing_properties)
 527
 528         prefix = "demo_"
 529         log_string(f"nb_{prefix}samples {n_epoch} {acc_nb_results}")
 530         log_string(
 531             f"property_{prefix}nb {n_epoch} requested {sum(acc_nb_requested_properties)} missing {sum(acc_nb_missing_properties)}"
 532         )
 533         log_string(
 534             f"property_{prefix}miss {n_epoch} {100*nb_missing_properties/nb_requested_properties:.02f}%"
 535         )
 536
 537         img = picoclvr.descr2img(result_descr, height=self.height, width=self.width)
 538
 539         if img.dim() == 5:
 540             if img.size(1) == 1:
 541                 img = F.pad(img.squeeze(1), pad=(1, 1, 1, 1), value=64)
 542             else:
 543                 img = torch.cat(
 544                     [
 545                         torchvision.utils.make_grid(x, padding=1, pad_value=64)[None]
 546                         for x in img
 547                     ],
 548                     0,
 549                 )
 550
 551         image_name = os.path.join(args.result_dir, f"picoclvr_result_{n_epoch:04d}.png")
 552         torchvision.utils.save_image(
 553             img / 255.0, image_name, nrow=nb_per_primer, padding=1, pad_value=0.0
 554         )
 555         log_string(f"wrote {image_name}")
 556
 557
 558 ######################################################################
 559
 560
 561 class TaskMNIST(Task):
 562     def __init__(self, batch_size, device=torch.device("cpu")):
 563         self.device = device
 564         self.batch_size = batch_size
 565
 566     def batches(self, split="train"):
 567         assert split in {"train", "test"}
 568         data_set = torchvision.datasets.MNIST(
 569             root="./data", train=(split == "train"), download=True
 570         )
 571         data_input = data_set.data.view(-1, 28 * 28).long()
 572         if args.nb_train_samples is not None:
 573             data_input = data_input[: args.nb_train_samples]
 574         for batch in tqdm.tqdm(
 575             data_input.split(self.batch_size), desc=f"epoch-{split}"
 576         ):
 577             yield batch
 578
 579     def vocabulary_size(self):
 580         return 256
 581
 582     def produce_results(self, n_epoch, model):
 583         results = torch.empty(64, 28 * 28, device=self.device, dtype=torch.int64)
 584         ar_mask = torch.full_like(results, 1)
 585         masked_inplace_autoregression(
 586             model, self.batch_size, results, ar_mask, device=self.device
 587         )
 588         image_name = os.path.join(args.result_dir, f"mnist_result_{n_epoch:04d}.png")
 589         torchvision.utils.save_image(
 590             1 - results.reshape(-1, 1, 28, 28) / 255.0,
 591             image_name,
 592             nrow=16,
 593             pad_value=0.8,
 594         )
 595         log_string(f"wrote {image_name}")
 596
 597
 598 ######################################################################
 599
 600 import maze
 601
 602
 603 class TaskMaze(Task):
 604     def map2seq(self, *m):
 605         return torch.cat([x.flatten(1) for x in m], 1)
 606
 607     def seq2map(self, s):
 608         s = s.reshape(s.size(0), -1, self.height, self.width)
 609         return (s[:, k] for k in range(s.size(1)))
 610
 611     def __init__(
 612         self,
 613         nb_train_samples,
 614         nb_test_samples,
 615         batch_size,
 616         height,
 617         width,
 618         nb_walls,
 619         device=torch.device("cpu"),
 620     ):
 621         self.batch_size = batch_size
 622         self.height = height
 623         self.width = width
 624         self.device = device
 625
 626         train_mazes, train_paths, _ = maze.create_maze_data(
 627             nb_train_samples,
 628             height=height,
 629             width=width,
 630             nb_walls=nb_walls,
 631             progress_bar=lambda x: tqdm.tqdm(x, dynamic_ncols=True, desc=f"data-train"),
 632         )
 633         self.train_input = self.map2seq(train_mazes.to(device), train_paths.to(device))
 634
 635         test_mazes, test_paths, _ = maze.create_maze_data(
 636             nb_test_samples,
 637             height=height,
 638             width=width,
 639             nb_walls=nb_walls,
 640             progress_bar=lambda x: tqdm.tqdm(x, dynamic_ncols=True, desc=f"data-test"),
 641         )
 642         self.test_input = self.map2seq(test_mazes.to(device), test_paths.to(device))
 643
 644         self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1
 645
 646     def batches(self, split="train", nb_to_use=-1, desc=None):
 647         assert split in {"train", "test"}
 648         input = self.train_input if split == "train" else self.test_input
 649         if nb_to_use > 0:
 650             input = input[:nb_to_use]
 651         if desc is None:
 652             desc = f"epoch-{split}"
 653         for batch in tqdm.tqdm(
 654             input.split(self.batch_size), dynamic_ncols=True, desc=desc
 655         ):
 656             yield batch
 657
 658     def vocabulary_size(self):
 659         return self.nb_codes
 660
 661     def compute_error(self, model, split="train", nb_to_use=-1):
 662         nb_total, nb_correct = 0, 0
 663         count = torch.zeros(
 664             self.width * self.height,
 665             self.width * self.height,
 666             device=self.device,
 667             dtype=torch.int64,
 668         )
 669         for input in tqdm.tqdm(
 670             task.batches(split, nb_to_use),
 671             dynamic_ncols=True,
 672             desc=f"test-mazes",
 673         ):
 674             result = input.clone()
 675             ar_mask = result.new_zeros(result.size())
 676             ar_mask[:, self.height * self.width :] = 1
 677             result *= 1 - ar_mask
 678             masked_inplace_autoregression(
 679                 model,
 680                 self.batch_size,
 681                 result,
 682                 ar_mask,
 683                 progress_bar_desc=None,
 684                 device=self.device,
 685             )
 686             mazes, paths = self.seq2map(result)
 687             path_correctness = maze.path_correctness(mazes, paths)
 688             nb_correct += path_correctness.long().sum()
 689             nb_total += mazes.size(0)
 690
 691             optimal_path_lengths = (
 692                 (input[:, self.height * self.width :] == maze.v_path).long().sum(1)
 693             )
 694             predicted_path_lengths = (
 695                 (result[:, self.height * self.width :] == maze.v_path).long().sum(1)
 696             )
 697             optimal_path_lengths = optimal_path_lengths[path_correctness]
 698             predicted_path_lengths = predicted_path_lengths[path_correctness]
 699             count[optimal_path_lengths, predicted_path_lengths] += 1
 700
 701         if count.max() == 0:
 702             count = None
 703         else:
 704             count = count[
 705                 : count.sum(1).nonzero().max() + 1, : count.sum(0).nonzero().max() + 1
 706             ]
 707
 708         return nb_total, nb_correct, count
 709
 710     def produce_results(self, n_epoch, model):
 711         with torch.autograd.no_grad():
 712             t = model.training
 713             model.eval()
 714
 715             train_nb_total, train_nb_correct, count = self.compute_error(
 716                 model, "train", nb_to_use=1000
 717             )
 718             log_string(
 719                 f"accuracy_train {n_epoch} nb_total {train_nb_total} nb_correct {train_nb_correct} accuracy {(100.0*train_nb_correct)/train_nb_total:.02f}%"
 720             )
 721
 722             test_nb_total, test_nb_correct, count = self.compute_error(
 723                 model, "test", nb_to_use=1000
 724             )
 725             log_string(
 726                 f"accuracy_test {n_epoch} nb_total {test_nb_total} nb_correct {test_nb_correct} accuracy {(100.0*test_nb_correct)/test_nb_total:.02f}%"
 727             )
 728
 729             if count is not None:
 730                 proportion_optimal = count.diagonal().sum().float() / count.sum()
 731                 log_string(f"proportion_optimal_test {proportion_optimal*100:.02f}%")
 732                 with open(
 733                     os.path.join(args.result_dir, f"maze_result_{n_epoch:04d}.txt"), "w"
 734                 ) as f:
 735                     for i in range(count.size(0)):
 736                         for j in range(count.size(1)):
 737                             eol = " " if j < count.size(1) - 1 else "\n"
 738                             f.write(f"{count[i,j]}{eol}")
 739
 740             input = self.test_input[:48]
 741             result = input.clone()
 742             ar_mask = result.new_zeros(result.size())
 743             ar_mask[:, self.height * self.width :] = 1
 744             result *= 1 - ar_mask
 745             masked_inplace_autoregression(
 746                 model, self.batch_size, result, ar_mask, device=self.device
 747             )
 748
 749             mazes, paths = self.seq2map(input)
 750             _, predicted_paths = self.seq2map(result)
 751
 752             filename = os.path.join(args.result_dir, f"maze_result_{n_epoch:04d}.png")
 753             maze.save_image(
 754                 filename,
 755                 mazes=mazes,
 756                 target_paths=paths,
 757                 predicted_paths=predicted_paths,
 758                 path_correct=maze.path_correctness(mazes, predicted_paths),
 759                 path_optimal=maze.path_optimality(paths, predicted_paths),
 760             )
 761             log_string(f"wrote {filename}")
 762
 763             model.train(t)
 764
 765
 766 ######################################################################
 767
 768
 769 import snake
 770
 771
 772 class TaskSnake(Task):
 773     def __init__(
 774         self,
 775         nb_train_samples,
 776         nb_test_samples,
 777         batch_size,
 778         height,
 779         width,
 780         nb_colors,
 781         length,
 782         prompt_length,
 783         device=torch.device("cpu"),
 784     ):
 785         self.batch_size = batch_size
 786         self.height = height
 787         self.width = width
 788         self.device = device
 789         self.prompt_length = prompt_length
 790
 791         self.train_input, self.train_prior_visits, _, _ = snake.generate_sequences(
 792             nb_train_samples,
 793             height,
 794             width,
 795             nb_colors,
 796             length,
 797             prompt_length,
 798             self.device,
 799         )
 800         self.test_input, self.test_prior_visits, _, _ = snake.generate_sequences(
 801             nb_test_samples,
 802             height,
 803             width,
 804             nb_colors,
 805             length,
 806             prompt_length,
 807             self.device,
 808         )
 809
 810         self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1
 811
 812     def batches(self, split="train", nb_to_use=-1, desc=None):
 813         assert split in {"train", "test"}
 814         input = self.train_input if split == "train" else self.test_input
 815         if nb_to_use > 0:
 816             input = input[:nb_to_use]
 817         if desc is None:
 818             desc = f"epoch-{split}"
 819         for batch in tqdm.tqdm(
 820             input.split(self.batch_size), dynamic_ncols=True, desc=desc
 821         ):
 822             yield batch
 823
 824     def vocabulary_size(self):
 825         return self.nb_codes
 826
 827     def produce_results(self, n_epoch, model):
 828         with torch.autograd.no_grad():
 829             t = model.training
 830             model.eval()
 831
 832             def compute_nb_correct(input, prior_visits):
 833                 result = input.clone()
 834                 i = torch.arange(result.size(1), device=result.device)[None, :]
 835                 ar_mask = (
 836                     torch.logical_and(i >= self.prompt_length * 2, i % 2 == 0)
 837                     .long()
 838                     .expand_as(result)
 839                 )
 840                 result *= 1 - ar_mask
 841
 842                 # snake.solver(result,ar_mask)
 843
 844                 masked_inplace_autoregression(
 845                     model, self.batch_size, result, ar_mask, device=self.device
 846                 )
 847
 848                 nb_total = ((prior_visits > 0) * ar_mask).sum()
 849
 850                 nb_correct = (
 851                     (result == input).long() * (prior_visits > 0) * ar_mask
 852                 ).sum()
 853
 854                 # nb_total = result.size(0)
 855                 # nb_correct = ((result - input).abs().sum(1) == 0).sum()
 856
 857                 return nb_total, nb_correct
 858
 859             # train_nb_total, train_nb_correct = compute_nb_correct(
 860             # self.train_input, self.train_prior_visits
 861             # )
 862
 863             # log_string(
 864             # f"accuracy_train nb_total {train_nb_total} nb_correct {train_nb_correct} accuracy {(100.0*train_nb_correct)/train_nb_total:.02f}%"
 865             # )
 866
 867             test_nb_total, test_nb_correct = compute_nb_correct(
 868                 self.test_input[:1000], self.test_prior_visits[:1000]
 869             )
 870
 871             log_string(
 872                 f"accuracy_test {n_epoch} nb_total {test_nb_total} nb_correct {test_nb_correct} accuracy {(100.0*test_nb_correct)/test_nb_total:.02f}%"
 873             )
 874
 875             model.train(t)
 876
 877
 878 ######################################################################
 879
 880
 881 import stack
 882
 883
 884 class TaskStack(Task):
 885     def __init__(
 886         self,
 887         nb_train_samples,
 888         nb_test_samples,
 889         batch_size,
 890         nb_steps,
 891         nb_stacks,
 892         nb_digits,
 893         fraction_values_for_train=None,
 894         device=torch.device("cpu"),
 895     ):
 896         self.batch_size = batch_size
 897         self.nb_steps = nb_steps
 898         self.nb_stacks = nb_stacks
 899         self.nb_digits = nb_digits
 900         self.device = device
 901
 902         if fraction_values_for_train is None:
 903             values_for_train = None
 904             values_for_test = None
 905         else:
 906             all = torch.randperm(10**nb_digits)
 907             nb_for_train = int(all.size(0) * fraction_values_for_train)
 908             values_for_train = all[:nb_for_train]
 909             values_for_test = all[nb_for_train:]
 910
 911         self.train_input, self.train_stack_counts = stack.generate_sequences(
 912             nb_train_samples,
 913             nb_steps,
 914             nb_stacks,
 915             nb_digits,
 916             values_for_train,
 917             self.device,
 918         )
 919
 920         self.test_input, self.test_stack_counts = stack.generate_sequences(
 921             nb_test_samples,
 922             nb_steps,
 923             nb_stacks,
 924             nb_digits,
 925             values_for_test,
 926             self.device,
 927         )
 928
 929         i = torch.logical_and(self.test_input % 2 == 1, self.test_input < 2 * nb_stacks)
 930         counts = self.test_stack_counts.flatten()[i.flatten()]
 931         counts = F.one_hot(counts).sum(0)
 932         log_string(f"test_pop_stack_counts {counts}")
 933
 934         self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1
 935
 936     def batches(self, split="train", nb_to_use=-1, desc=None):
 937         assert split in {"train", "test"}
 938         input = self.train_input if split == "train" else self.test_input
 939         if nb_to_use > 0:
 940             input = input[:nb_to_use]
 941         if desc is None:
 942             desc = f"epoch-{split}"
 943         for batch in tqdm.tqdm(
 944             input.split(self.batch_size), dynamic_ncols=True, desc=desc
 945         ):
 946             yield batch
 947
 948     def vocabulary_size(self):
 949         return self.nb_codes
 950
 951     def produce_results(self, n_epoch, model):
 952         with torch.autograd.no_grad():
 953             t = model.training
 954             model.eval()
 955
 956             def compute_nb_correct(input):
 957                 result = input.clone()
 958                 stack.remove_popped_values(result, self.nb_stacks, self.nb_digits)
 959                 ar_mask = (result != input).long()
 960                 masked_inplace_autoregression(
 961                     model, self.batch_size, result, ar_mask, device=self.device
 962                 )
 963
 964                 errors = ((result != input).long() * ar_mask).reshape(
 965                     -1, 1 + self.nb_digits
 966                 )
 967                 ar_mask = ar_mask.reshape(-1, 1 + self.nb_digits)
 968
 969                 nb_total = ar_mask.max(1).values.sum()
 970                 nb_correct = nb_total - errors.max(1).values.sum()
 971
 972                 return nb_total, nb_correct
 973
 974             test_nb_total, test_nb_correct = compute_nb_correct(self.test_input[:1000])
 975
 976             log_string(
 977                 f"accuracy_test {n_epoch} nb_total {test_nb_total} nb_correct {test_nb_correct} accuracy {(100.0*test_nb_correct)/test_nb_total:.02f}%"
 978             )
 979
 980             ##############################################################
 981             # Log a few generated sequences
 982             input = self.test_input[:10, : 12 * (1 + self.nb_digits)]
 983             result = input.clone()
 984             stack.remove_popped_values(result, self.nb_stacks, self.nb_digits)
 985             ar_mask = (result != input).long()
 986             for n in range(result.size(0)):
 987                 log_string(
 988                     f"test_before {stack.seq_to_str(result[n],nb_stacks=self.nb_stacks,nb_digits=self.nb_digits)}"
 989                 )
 990             masked_inplace_autoregression(
 991                 model, self.batch_size, result, ar_mask, device=self.device
 992             )
 993             for n in range(result.size(0)):
 994                 log_string(
 995                     f"test_after  {stack.seq_to_str(result[n],nb_stacks=self.nb_stacks,nb_digits=self.nb_digits)}"
 996                 )
 997             ##############################################################
 998
 999             model.train(t)
1000
1001
1002 ######################################################################
1003
1004
1005 import expr
1006
1007
1008 class TaskExpr(Task):
1009     def __init__(
1010         self,
1011         nb_train_samples,
1012         nb_test_samples,
1013         batch_size,
1014         device=torch.device("cpu"),
1015     ):
1016         self.batch_size = batch_size
1017         self.device = device
1018
1019         train_sequences = expr.generate_sequences(nb_train_samples)
1020         test_sequences = expr.generate_sequences(nb_test_samples)
1021         self.char2id = dict([ (c,n) for n,c in enumerate(set("".join(train_sequences + test_sequences))) ])
1022         self.id2char = dict([ (n,c) for n,c in self.char2id.items() ])
1023         len_max = max([len(x) for x in train_sequences + test_sequences])
1024         self.train_input = torch.cat([torch.tensor([char2id(c) for c in s + " "*(len_max-len(s))] for s in train_sequences)], 0)
1025         self.test_input = torch.cat([torch.tensor([char2id(c) for c in s + " "*(len_max-len(s))] for s in test_sequences)], 0)
1026         self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1
1027
1028     def batches(self, split="train", nb_to_use=-1, desc=None):
1029         assert split in {"train", "test"}
1030         input = self.train_input if split == "train" else self.test_input
1031         if nb_to_use > 0:
1032             input = input[:nb_to_use]
1033         if desc is None:
1034             desc = f"epoch-{split}"
1035         for batch in tqdm.tqdm(
1036             input.split(self.batch_size), dynamic_ncols=True, desc=desc
1037         ):
1038             yield batch
1039
1040     def vocabulary_size(self):
1041         return self.nb_codes
1042
1043     def produce_results(self, n_epoch, model):
1044         # !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
1045         with torch.autograd.no_grad():
1046             t = model.training
1047             model.eval()
1048
1049             def compute_nb_correct(input):
1050                 result = input.clone()
1051                 stack.remove_popped_values(result, self.nb_stacks, self.nb_digits)
1052                 ar_mask = (result != input).long()
1053                 masked_inplace_autoregression(
1054                     model, self.batch_size, result, ar_mask, device=self.device
1055                 )
1056
1057                 errors = ((result != input).long() * ar_mask).reshape(
1058                     -1, 1 + self.nb_digits
1059                 )
1060                 ar_mask = ar_mask.reshape(-1, 1 + self.nb_digits)
1061
1062                 nb_total = ar_mask.max(1).values.sum()
1063                 nb_correct = nb_total - errors.max(1).values.sum()
1064
1065                 return nb_total, nb_correct
1066
1067             test_nb_total, test_nb_correct = compute_nb_correct(self.test_input[:1000])
1068
1069             log_string(
1070                 f"accuracy_test {n_epoch} nb_total {test_nb_total} nb_correct {test_nb_correct} accuracy {(100.0*test_nb_correct)/test_nb_total:.02f}%"
1071             )
1072
1073             ##############################################################
1074             # Log a few generated sequences
1075             input = self.test_input[:10, : 12 * (1 + self.nb_digits)]
1076             result = input.clone()
1077             stack.remove_popped_values(result, self.nb_stacks, self.nb_digits)
1078             ar_mask = (result != input).long()
1079             for n in range(result.size(0)):
1080                 log_string(
1081                     f"test_before {stack.seq_to_str(result[n],nb_stacks=self.nb_stacks,nb_digits=self.nb_digits)}"
1082                 )
1083             masked_inplace_autoregression(
1084                 model, self.batch_size, result, ar_mask, device=self.device
1085             )
1086             for n in range(result.size(0)):
1087                 log_string(
1088                     f"test_after  {stack.seq_to_str(result[n],nb_stacks=self.nb_stacks,nb_digits=self.nb_digits)}"
1089                 )
1090             ##############################################################
1091
1092             model.train(t)
1093
1094
1095 ######################################################################
1096
1097
1098 def picoclvr_pruner_horizontal_green(p):
1099     return not ("green" in p and ("left" in p or "right" in p))
1100
1101
1102 picoclvr_pruner_train = (
1103     picoclvr_pruner_horizontal_green
1104     if args.picocvlr_prune_properties in {"train+eval"}
1105     else None
1106 )
1107
1108 picoclvr_pruner_eval = (
1109     (lambda p: not picoclvr_pruner_horizontal_green(p))
1110     if args.picocvlr_prune_properties in {"train+eval", "eval"}
1111     else None
1112 )
1113
1114 ######################################################################
1115
1116 if args.task == "picoclvr":
1117     task = TaskPicoCLVR(
1118         nb_train_samples=args.nb_train_samples,
1119         nb_test_samples=args.nb_test_samples,
1120         batch_size=args.batch_size,
1121         height=args.picoclvr_height,
1122         width=args.picoclvr_width,
1123         nb_colors=args.picoclvr_nb_colors,
1124         device=device,
1125         pruner_train=picoclvr_pruner_train,
1126         pruner_eval=picoclvr_pruner_eval,
1127     )
1128
1129 elif args.task == "mnist":
1130     task = TaskMNIST(
1131         batch_size=args.batch_size,
1132         device=device,
1133     )
1134
1135 elif args.task == "maze":
1136     task = TaskMaze(
1137         nb_train_samples=args.nb_train_samples,
1138         nb_test_samples=args.nb_test_samples,
1139         batch_size=args.batch_size,
1140         height=args.maze_height,
1141         width=args.maze_width,
1142         nb_walls=args.maze_nb_walls,
1143         device=device,
1144     )
1145
1146 elif args.task == "snake":
1147     task = TaskSnake(
1148         nb_train_samples=args.nb_train_samples,
1149         nb_test_samples=args.nb_test_samples,
1150         batch_size=args.batch_size,
1151         height=args.snake_height,
1152         width=args.snake_width,
1153         nb_colors=args.snake_nb_colors,
1154         length=args.snake_length,
1155         prompt_length=args.snake_length // 2,
1156         device=device,
1157     )
1158
1159 elif args.task == "stack":
1160     task = TaskStack(
1161         nb_train_samples=args.nb_train_samples,
1162         nb_test_samples=args.nb_test_samples,
1163         batch_size=args.batch_size,
1164         nb_steps=args.stack_nb_steps,
1165         nb_stacks=args.stack_nb_stacks,
1166         nb_digits=args.stack_nb_digits,
1167         fraction_values_for_train=args.stack_fraction_values_for_train,
1168         device=device,
1169     )
1170
1171 elif args.task == "expr":
1172     task = TaskExpr(
1173         nb_train_samples=args.nb_train_samples,
1174         nb_test_samples=args.nb_test_samples,
1175         batch_size=args.batch_size,
1176         device=device,
1177     )
1178
1179 else:
1180     raise ValueError(f"Unknown task {args.task}")
1181
1182 ######################################################################
1183
1184 log_string(f"device {device}")
1185
1186 vocabulary_size = task.vocabulary_size()
1187
1188 log_string(f"vocabulary_size {vocabulary_size}")
1189
1190 ##############################
1191
1192 model = mygpt.MyGPT(
1193     vocabulary_size=vocabulary_size,
1194     dim_model=args.dim_model,
1195     dim_keys=args.dim_keys,
1196     dim_hidden=args.dim_hidden,
1197     nb_heads=args.nb_heads,
1198     nb_blocks=args.nb_blocks,
1199     causal=True,
1200     dropout=args.dropout,
1201 )
1202
1203 model.to(device)
1204
1205 nb_parameters = sum(p.numel() for p in model.parameters())
1206 log_string(f"nb_parameters {nb_parameters} ({int(nb_parameters/1e6)}M)")
1207
1208 ######################################################################
1209
1210 nb_epochs_finished = 0
1211
1212 if args.no_checkpoint:
1213     log_string(f"not trying to load checkpoint.")
1214
1215 else:
1216     try:
1217         checkpoint_name = os.path.join(args.result_dir, args.checkpoint_name)
1218         checkpoint = torch.load(checkpoint_name)
1219         nb_epochs_finished = checkpoint["nb_epochs_finished"]
1220         model.load_state_dict(checkpoint["model_state"])
1221         torch.set_rng_state(checkpoint["rng_state"])
1222         if torch.cuda.is_available():
1223             torch.cuda.set_rng_state(checkpoint["cuda_rng_state"])
1224
1225         log_string(f"checkpoint loaded with {nb_epochs_finished} epochs finished.")
1226
1227     except FileNotFoundError:
1228         log_string("starting from scratch.")
1229
1230     except:
1231         log_string("error when loading the checkpoint.")
1232         exit(1)
1233
1234 ######################################################################
1235
1236 nb_epochs = args.nb_epochs if args.nb_epochs > 0 else nb_epochs_default
1237
1238 token_count = 0
1239 for input in task.batches(split="train"):
1240     token_count += F.one_hot(input, num_classes=task.vocabulary_size()).sum((0, 1))
1241 token_probas = token_count / token_count.sum()
1242 entropy = -torch.xlogy(token_probas, token_probas).sum()
1243 train_set_perplexity = math.exp(entropy)
1244
1245 ##############################
1246
1247 if args.learning_rate_schedule == "cos":
1248     learning_rate_schedule = {}
1249     for n_epoch in range(args.nb_epochs):
1250         u = n_epoch / args.nb_epochs * math.pi
1251         learning_rate_schedule[n_epoch] = args.learning_rate * 0.5 * (1 + math.cos(u))
1252 else:
1253     u = {
1254         int(k): float(v)
1255         for k, v in [
1256             tuple(x.split(":")) for x in args.learning_rate_schedule.split(",")
1257         ]
1258     }
1259
1260     learning_rate_schedule = {}
1261     learning_rate = args.learning_rate
1262     for n_epoch in range(args.nb_epochs):
1263         if n_epoch in u:
1264             learning_rate = u[n_epoch]
1265         learning_rate_schedule[n_epoch] = learning_rate
1266
1267 log_string(f"learning_rate_schedule {learning_rate_schedule}")
1268
1269 ##############################
1270
1271 nb_samples_seen = 0
1272
1273 if nb_epochs_finished >= nb_epochs:
1274     task.produce_results(nb_epochs_finished, model)
1275
1276 for n_epoch in range(nb_epochs_finished, nb_epochs):
1277     learning_rate = learning_rate_schedule[n_epoch]
1278
1279     log_string(f"learning_rate {learning_rate}")
1280
1281     if args.optim == "sgd":
1282         optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
1283     elif args.optim == "adam":
1284         optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
1285     elif args.optim == "adamw":
1286         optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
1287     else:
1288         raise ValueError(f"Unknown optimizer {args.optim}.")
1289
1290     model.train()
1291
1292     nb_train_samples, acc_train_loss = 0, 0.0
1293
1294     for input in task.batches(split="train"):
1295         input = input.to(device)
1296         output = model(mygpt.BracketedSequence(input)).x
1297         loss = F.cross_entropy(output.transpose(1, 2), input)
1298         acc_train_loss += loss.item() * input.size(0)
1299         nb_train_samples += input.size(0)
1300         nb_samples_seen += input.size(0)
1301
1302         optimizer.zero_grad()
1303         loss.backward()
1304         optimizer.step()
1305
1306     with torch.autograd.no_grad():
1307         model.eval()
1308
1309         nb_test_samples, acc_test_loss = 0, 0.0
1310
1311         for input in task.batches(split="test"):
1312             input = input.to(device)
1313
1314             output = model(mygpt.BracketedSequence(input)).x
1315             loss = F.cross_entropy(output.transpose(1, 2), input)
1316             acc_test_loss += loss.item() * input.size(0)
1317             nb_test_samples += input.size(0)
1318
1319         train_perplexity = math.exp(min(100, acc_train_loss / nb_train_samples))
1320         test_perplexity = math.exp(min(100, acc_test_loss / nb_test_samples))
1321
1322         log_string(
1323             f"perplexity {n_epoch} train_set {train_set_perplexity} train_prediction {train_perplexity} test_prediction {test_perplexity}"
1324         )
1325
1326         task.produce_results(n_epoch, model)
1327
1328     checkpoint = {
1329         "nb_epochs_finished": n_epoch + 1,
1330         "model_state": model.state_dict(),
1331         "rng_state": torch.get_rng_state(),
1332     }
1333
1334     if torch.cuda.is_available():
1335         checkpoint["cuda_rng_state"] = torch.cuda.get_rng_state()
1336
1337     checkpoint_name = os.path.join(args.result_dir, args.checkpoint_name)
1338     torch.save(checkpoint, checkpoint_name)
1339     log_string(f"saved checkpoint {checkpoint_name}")
1340
1341 ######################################################################