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