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