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=1.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, self.width * self.height, device=self.device, dtype=torch.int64
 626         )
 627         for input in task.batches(split, nb_to_use):
 628             result = input.clone()
 629             ar_mask = result.new_zeros(result.size())
 630             ar_mask[:, self.height * self.width :] = 1
 631             result *= 1 - ar_mask
 632             masked_inplace_autoregression(
 633                 model, self.batch_size, result, ar_mask, device=self.device
 634             )
 635             mazes, paths = self.seq2map(result)
 636             path_correctness = maze.path_correctness(mazes, paths)
 637             nb_correct += path_correctness.long().sum()
 638             nb_total += mazes.size(0)
 639
 640             optimal_path_lengths = (
 641                 (input[:, self.height * self.width :] == maze.v_path).long().sum(1)
 642             )
 643             predicted_path_lengths = (
 644                 (result[:, self.height * self.width :] == maze.v_path).long().sum(1)
 645             )
 646             optimal_path_lengths = optimal_path_lengths[path_correctness]
 647             predicted_path_lengths = predicted_path_lengths[path_correctness]
 648             count[optimal_path_lengths, predicted_path_lengths] += 1
 649
 650         if count.max() == 0:
 651             count = None
 652         else:
 653             count = count[
 654                 : count.sum(1).nonzero().max() + 1, : count.sum(0).nonzero().max() + 1
 655             ]
 656
 657         return nb_total, nb_correct, count
 658
 659     def produce_results(self, n_epoch, model):
 660         with torch.autograd.no_grad():
 661             t = model.training
 662             model.eval()
 663
 664             train_nb_total, train_nb_correct, count = self.compute_error(
 665                 model, "train", nb_to_use=1000
 666             )
 667             log_string(
 668                 f"accuracy_train nb_total {train_nb_total} nb_correct {train_nb_correct} accuracy {(100.0*train_nb_correct)/train_nb_total:.02f}%"
 669             )
 670
 671             test_nb_total, test_nb_correct, count = self.compute_error(
 672                 model, "test", nb_to_use=1000
 673             )
 674             log_string(
 675                 f"accuracy_test nb_total {test_nb_total} nb_correct {test_nb_correct} accuracy {(100.0*test_nb_correct)/test_nb_total:.02f}%"
 676             )
 677
 678             if count is not None:
 679                 proportion_optimal = count.diagonal().sum().float() / count.sum()
 680                 log_string(f"proportion_optimal_test {proportion_optimal*100:.02f}%")
 681                 with open(
 682                     os.path.join(args.result_dir, f"maze_result_{n_epoch:04d}.txt"), "w"
 683                 ) as f:
 684                     for i in range(count.size(0)):
 685                         for j in range(count.size(1)):
 686                             eol = " " if j < count.size(1) - 1 else "\n"
 687                             f.write(f"{count[i,j]}{eol}")
 688
 689             input = self.test_input[:48]
 690             result = input.clone()
 691             ar_mask = result.new_zeros(result.size())
 692             ar_mask[:, self.height * self.width :] = 1
 693             result *= 1 - ar_mask
 694             masked_inplace_autoregression(
 695                 model, self.batch_size, result, ar_mask, device=self.device
 696             )
 697
 698             mazes, paths = self.seq2map(input)
 699             _, predicted_paths = self.seq2map(result)
 700
 701             filename = os.path.join(args.result_dir, f"maze_result_{n_epoch:04d}.png")
 702             maze.save_image(
 703                 filename,
 704                 mazes=mazes,
 705                 target_paths=paths,
 706                 predicted_paths=predicted_paths,
 707                 path_correct=maze.path_correctness(mazes, predicted_paths),
 708             )
 709             log_string(f"wrote {filename}")
 710
 711             model.train(t)
 712
 713
 714 ######################################################################
 715
 716
 717 import snake
 718
 719
 720 class TaskSnake(Task):
 721     def __init__(
 722         self,
 723         nb_train_samples,
 724         nb_test_samples,
 725         batch_size,
 726         height,
 727         width,
 728         nb_colors,
 729         length,
 730         prompt_length,
 731         device=torch.device("cpu"),
 732     ):
 733         self.batch_size = batch_size
 734         self.height = height
 735         self.width = width
 736         self.device = device
 737         self.prompt_length = prompt_length
 738
 739         self.train_input, self.train_prior_visits, _, _ = snake.generate_sequences(
 740             nb_train_samples,
 741             height,
 742             width,
 743             nb_colors,
 744             length,
 745             prompt_length,
 746             self.device,
 747         )
 748         self.test_input, self.test_prior_visits, _, _ = snake.generate_sequences(
 749             nb_test_samples,
 750             height,
 751             width,
 752             nb_colors,
 753             length,
 754             prompt_length,
 755             self.device,
 756         )
 757
 758         self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1
 759
 760     def batches(self, split="train", nb_to_use=-1, desc=None):
 761         assert split in {"train", "test"}
 762         input = self.train_input if split == "train" else self.test_input
 763         if nb_to_use > 0:
 764             input = input[:nb_to_use]
 765         if desc is None:
 766             desc = f"epoch-{split}"
 767         for batch in tqdm.tqdm(
 768             input.split(self.batch_size), dynamic_ncols=True, desc=desc
 769         ):
 770             yield batch
 771
 772     def vocabulary_size(self):
 773         return self.nb_codes
 774
 775     def produce_results(self, n_epoch, model):
 776         with torch.autograd.no_grad():
 777             t = model.training
 778             model.eval()
 779
 780             def compute_nb_correct(input, prior_visits):
 781                 result = input.clone()
 782                 i = torch.arange(result.size(1), device=result.device)[None, :]
 783                 ar_mask = (
 784                     torch.logical_and(i >= self.prompt_length * 2, i % 2 == 0)
 785                     .long()
 786                     .expand_as(result)
 787                 )
 788                 result *= 1 - ar_mask
 789
 790                 # snake.solver(result,ar_mask)
 791
 792                 masked_inplace_autoregression(
 793                     model, self.batch_size, result, ar_mask, device=self.device
 794                 )
 795
 796                 nb_total = ((prior_visits > 0) * ar_mask).sum()
 797
 798                 nb_correct = (
 799                     (result == input).long() * (prior_visits > 0) * ar_mask
 800                 ).sum()
 801
 802                 # nb_total = result.size(0)
 803                 # nb_correct = ((result - input).abs().sum(1) == 0).sum()
 804
 805                 return nb_total, nb_correct
 806
 807             # train_nb_total, train_nb_correct = compute_nb_correct(
 808             # self.train_input, self.train_prior_visits
 809             # )
 810
 811             # log_string(
 812             # f"accuracy_train nb_total {train_nb_total} nb_correct {train_nb_correct} accuracy {(100.0*train_nb_correct)/train_nb_total:.02f}%"
 813             # )
 814
 815             test_nb_total, test_nb_correct = compute_nb_correct(
 816                 self.test_input[:1000], self.test_prior_visits[:1000]
 817             )
 818
 819             log_string(
 820                 f"accuracy_test nb_total {test_nb_total} nb_correct {test_nb_correct} accuracy {(100.0*test_nb_correct)/test_nb_total:.02f}%"
 821             )
 822
 823             model.train(t)
 824
 825
 826 ######################################################################
 827
 828
 829 def picoclvr_pruner_horizontal_green(p):
 830     return not ("green" in p and ("left" in p or "right" in p))
 831
 832
 833 picoclvr_pruner_train = (
 834     picoclvr_pruner_horizontal_green
 835     if args.picocvlr_prune_properties in {"train+eval"}
 836     else None
 837 )
 838
 839 picoclvr_pruner_eval = (
 840     (lambda p: not picoclvr_pruner_horizontal_green(p))
 841     if args.picocvlr_prune_properties in {"train+eval", "eval"}
 842     else None
 843 )
 844
 845 ######################################################################
 846
 847 if args.task == "picoclvr":
 848     task = TaskPicoCLVR(
 849         nb_train_samples=args.nb_train_samples,
 850         nb_test_samples=args.nb_test_samples,
 851         batch_size=args.batch_size,
 852         height=args.picoclvr_height,
 853         width=args.picoclvr_width,
 854         nb_colors=args.picoclvr_nb_colors,
 855         device=device,
 856         pruner_train=picoclvr_pruner_train,
 857         pruner_eval=picoclvr_pruner_eval,
 858     )
 859
 860 elif args.task == "mnist":
 861     task = TaskMNIST(
 862         batch_size=args.batch_size,
 863         device=device,
 864     )
 865
 866 elif args.task == "maze":
 867     task = TaskMaze(
 868         nb_train_samples=args.nb_train_samples,
 869         nb_test_samples=args.nb_test_samples,
 870         batch_size=args.batch_size,
 871         height=args.maze_height,
 872         width=args.maze_width,
 873         nb_walls=args.maze_nb_walls,
 874         device=device,
 875     )
 876
 877 elif args.task == "snake":
 878     task = TaskSnake(
 879         nb_train_samples=args.nb_train_samples,
 880         nb_test_samples=args.nb_test_samples,
 881         batch_size=args.batch_size,
 882         height=args.snake_height,
 883         width=args.snake_width,
 884         nb_colors=args.snake_nb_colors,
 885         length=args.snake_length,
 886         prompt_length=args.snake_length // 2,
 887         device=device,
 888     )
 889
 890 else:
 891     raise ValueError(f"Unknown task {args.task}")
 892
 893 ######################################################################
 894
 895 log_string(f"device {device}")
 896
 897 vocabulary_size = task.vocabulary_size()
 898
 899 log_string(f"vocabulary_size {vocabulary_size}")
 900
 901 ##############################
 902
 903 model = mygpt.MyGPT(
 904     vocabulary_size=vocabulary_size,
 905     dim_model=args.dim_model,
 906     dim_keys=args.dim_keys,
 907     dim_hidden=args.dim_hidden,
 908     nb_heads=args.nb_heads,
 909     nb_blocks=args.nb_blocks,
 910     causal=True,
 911     dropout=args.dropout,
 912 )
 913
 914 model.to(device)
 915
 916 nb_parameters = sum(p.numel() for p in model.parameters())
 917 log_string(f"nb_parameters {nb_parameters} ({int(nb_parameters/1e6)}M)")
 918
 919 ######################################################################
 920
 921 nb_epochs_finished = 0
 922
 923 if args.no_checkpoint:
 924     log_string(f"not trying to load checkpoint.")
 925
 926 else:
 927     try:
 928         checkpoint_name = os.path.join(args.result_dir, args.checkpoint_name)
 929         checkpoint = torch.load(checkpoint_name)
 930         nb_epochs_finished = checkpoint["nb_epochs_finished"]
 931         model.load_state_dict(checkpoint["model_state"])
 932         torch.set_rng_state(checkpoint["rng_state"])
 933         if torch.cuda.is_available():
 934             torch.cuda.set_rng_state(checkpoint["cuda_rng_state"])
 935
 936         log_string(f"checkpoint loaded with {nb_epochs_finished} epochs finished.")
 937
 938     except FileNotFoundError:
 939         log_string("starting from scratch.")
 940
 941     except:
 942         log_string("error when loading the checkpoint.")
 943         exit(1)
 944
 945 ######################################################################
 946
 947 nb_epochs = args.nb_epochs if args.nb_epochs > 0 else nb_epochs_default
 948
 949 token_count = 0
 950 for input in task.batches(split="train"):
 951     token_count += F.one_hot(input, num_classes=task.vocabulary_size()).sum((0, 1))
 952 token_probas = token_count / token_count.sum()
 953 entropy = -torch.xlogy(token_probas, token_probas).sum()
 954 train_set_perplexity = math.exp(entropy)
 955
 956 ##############################
 957
 958 if args.learning_rate_schedule == "cos":
 959     learning_rate_schedule = {}
 960     for n_epoch in range(args.nb_epochs):
 961         u = n_epoch / args.nb_epochs * math.pi
 962         learning_rate_schedule[n_epoch] = args.learning_rate * 0.5 * (1 + math.cos(u))
 963 else:
 964     u = {
 965         int(k): float(v)
 966         for k, v in [
 967             tuple(x.split(":")) for x in args.learning_rate_schedule.split(",")
 968         ]
 969     }
 970
 971     learning_rate_schedule = {}
 972     learning_rate = args.learning_rate
 973     for n_epoch in range(args.nb_epochs):
 974         if n_epoch in u:
 975             learning_rate = u[n_epoch]
 976         learning_rate_schedule[n_epoch] = learning_rate
 977
 978 log_string(f"learning_rate_schedule {learning_rate_schedule}")
 979
 980 ##############################
 981
 982 nb_samples_seen = 0
 983
 984 if nb_epochs_finished >= nb_epochs:
 985     task.produce_results(nb_epochs_finished, model)
 986
 987 for n_epoch in range(nb_epochs_finished, nb_epochs):
 988     learning_rate = learning_rate_schedule[n_epoch]
 989
 990     log_string(f"learning_rate {learning_rate}")
 991
 992     if args.optim == "sgd":
 993         optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
 994     elif args.optim == "adam":
 995         optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
 996     elif args.optim == "adamw":
 997         optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
 998     else:
 999         raise ValueError(f"Unknown optimizer {args.optim}.")
1000
1001     model.train()
1002
1003     nb_train_samples, acc_train_loss = 0, 0.0
1004
1005     for input in task.batches(split="train"):
1006         input = input.to(device)
1007         output = model(mygpt.BracketedSequence(input)).x
1008         loss = F.cross_entropy(output.transpose(1, 2), input)
1009         acc_train_loss += loss.item() * input.size(0)
1010         nb_train_samples += input.size(0)
1011         nb_samples_seen += input.size(0)
1012
1013         optimizer.zero_grad()
1014         loss.backward()
1015         optimizer.step()
1016
1017     with torch.autograd.no_grad():
1018         model.eval()
1019
1020         nb_test_samples, acc_test_loss = 0, 0.0
1021
1022         for input in task.batches(split="test"):
1023             input = input.to(device)
1024
1025             output = model(mygpt.BracketedSequence(input)).x
1026             loss = F.cross_entropy(output.transpose(1, 2), input)
1027             acc_test_loss += loss.item() * input.size(0)
1028             nb_test_samples += input.size(0)
1029
1030         train_perplexity = math.exp(min(100, acc_train_loss / nb_train_samples))
1031         test_perplexity = math.exp(min(100, acc_test_loss / nb_test_samples))
1032
1033         log_string(
1034             f"perplexity {n_epoch} train_set {train_set_perplexity} train_prediction {train_perplexity} test_prediction {test_perplexity}"
1035         )
1036
1037         task.produce_results(n_epoch, model)
1038
1039     checkpoint = {
1040         "nb_epochs_finished": n_epoch + 1,
1041         "model_state": model.state_dict(),
1042         "rng_state": torch.get_rng_state(),
1043     }
1044
1045     if torch.cuda.is_available():
1046         checkpoint["cuda_rng_state"] = torch.cuda.get_rng_state()
1047
1048     checkpoint_name = os.path.join(args.result_dir, args.checkpoint_name)
1049     torch.save(checkpoint, checkpoint_name)
1050     log_string(f"saved checkpoint {checkpoint_name}")
1051
1052 ######################################################################