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