tasks.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 import math, os, tqdm, warnings
   9
  10 import torch, torchvision
  11
  12 from torch import nn
  13 from torch.nn import functional as F
  14
  15 from mygpt import BracketedSequence
  16
  17 ######################################################################
  18
  19
  20 def masked_inplace_autoregression(
  21     model,
  22     batch_size,
  23     input,
  24     ar_mask,
  25     deterministic_synthesis,
  26     forbidden_tokens=None,
  27     logit_biases=None,
  28     progress_bar_desc="autoregression",
  29     device=torch.device("cpu"),
  30 ):
  31     assert input.size() == ar_mask.size()
  32
  33     batches = zip(input.split(batch_size), ar_mask.split(batch_size))
  34
  35     if progress_bar_desc is not None:
  36         batches = tqdm.tqdm(
  37             batches,
  38             dynamic_ncols=True,
  39             desc=progress_bar_desc,
  40             total=(input.size(0) + batch_size - 1) // batch_size,
  41         )
  42
  43     with torch.autograd.no_grad():
  44         t = model.training
  45         model.eval()
  46
  47         for input, ar_mask in batches:
  48             model.masked_inplace_autoregression(
  49                 input,
  50                 ar_mask,
  51                 deterministic_synthesis,
  52                 forbidden_tokens,
  53                 logit_biases,
  54             )
  55
  56         model.train(t)
  57
  58
  59 ######################################################################
  60
  61
  62 class Task:
  63     def batches(self, split="train", nb_to_use=-1, desc=None):
  64         pass
  65
  66     def vocabulary_size(self):
  67         pass
  68
  69     def produce_results(
  70         self, n_epoch, model, result_dir, logger, deterministic_synthesis
  71     ):
  72         pass
  73
  74
  75 ######################################################################
  76
  77 import world
  78
  79
  80 class World(Task):
  81     def save_image(self, input, result_dir, filename, logger):
  82         img = world.sample2img(input.to("cpu"), self.height, self.width)
  83         image_name = os.path.join(result_dir, filename)
  84         torchvision.utils.save_image(img.float() / 255.0, image_name, nrow=8, padding=2)
  85         logger(f"wrote {image_name}")
  86
  87     def make_ar_mask(self, input):
  88         b = torch.arange(input.size(1), device=input.device) > input.size(1) // 2
  89         return b.long()[None, :].expand_as(input)
  90
  91     def __init__(
  92         self,
  93         nb_train_samples,
  94         nb_test_samples,
  95         batch_size,
  96         result_dir=None,
  97         logger=None,
  98         device=torch.device("cpu"),
  99     ):
 100         super().__init__()
 101
 102         self.batch_size = batch_size
 103         self.device = device
 104         self.height = 6
 105         self.width = 8
 106
 107         self.train_input = world.generate(
 108             nb_train_samples, height=self.height, width=self.width
 109         ).to(device)
 110
 111         self.test_input = world.generate(
 112             nb_test_samples, height=self.height, width=self.width
 113         ).to(device)
 114
 115         self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1
 116
 117         self.train_quizzes = []
 118         self.test_quizzes = []
 119
 120         if result_dir is not None:
 121             self.save_image(
 122                 self.train_input[:96], result_dir, f"world_train.png", logger
 123             )
 124
 125     def batches(self, split="train", desc=None):
 126         assert split in {"train", "test"}
 127         if split == "train":
 128             input = self.train_input
 129             quizzes = self.train_quizzes
 130         else:
 131             input = self.test_input
 132             quizzes = self.test_quizzes
 133
 134         if len(quizzes) > 0:
 135             quizzes = torch.cat(quizzes, dim=0)
 136             if quizzes.size(0) > input.size(0) // 2:
 137                 i = torch.randperm(input.size(0))[: input.size(0) // 2]
 138                 quizzes = quizzes[i]
 139
 140             i = torch.randperm(input.size(0))[: input.size(0) - quizzes.size(0)]
 141             input = input[i]
 142
 143             self.nb_batch_samples_world = input.size(0)
 144             self.nb_batch_samples_quizzes = quizzes.size(0)
 145
 146             input = torch.cat([input, quizzes], dim=0)
 147         else:
 148             self.nb_batch_samples_world = input.size(0)
 149             self.nb_batch_samples_quizzes = 0
 150
 151         if desc is None:
 152             desc = f"epoch-{split}"
 153         for batch in tqdm.tqdm(
 154             input.split(self.batch_size), dynamic_ncols=True, desc=desc
 155         ):
 156             yield batch
 157
 158     def vocabulary_size(self):
 159         return self.nb_codes
 160
 161     def produce_results(
 162         self, n_epoch, model, result_dir, logger, deterministic_synthesis, nmax=1000
 163     ):
 164         def compute_accuracy(input, logger=None):
 165             input = input[:nmax]
 166             ar_mask = self.make_ar_mask(input)
 167             result = input.clone() * (1 - ar_mask)
 168
 169             masked_inplace_autoregression(
 170                 model,
 171                 self.batch_size,
 172                 result,
 173                 ar_mask,
 174                 deterministic_synthesis,
 175                 progress_bar_desc=None,
 176                 device=self.device,
 177             )
 178
 179             nb_total, nb_correct = (
 180                 input.size(0),
 181                 (input == result).long().min(dim=1).values.sum(),
 182             )
 183
 184             return nb_total, nb_correct
 185
 186         train_nb_total, train_nb_correct = compute_accuracy(self.train_input)
 187
 188         logger(
 189             f"accuracy_train {n_epoch} nb_total {train_nb_total} nb_correct {train_nb_correct} accuracy {(100.0*train_nb_correct)/train_nb_total:.02f}%"
 190         )
 191
 192         test_nb_total, test_nb_correct = compute_accuracy(self.test_input, logger)
 193
 194         logger(
 195             f"accuracy_test {n_epoch} nb_total {test_nb_total} nb_correct {test_nb_correct} accuracy {(100.0*test_nb_correct)/test_nb_total:.02f}%"
 196         )
 197
 198         main_test_accuracy = test_nb_correct / test_nb_total
 199         logger(f"main_test_accuracy {n_epoch} {main_test_accuracy}")
 200
 201         ##############################
 202
 203         input = self.test_input[:96]
 204         ar_mask = self.make_ar_mask(input)
 205         result = input.clone() * (1 - ar_mask)
 206
 207         masked_inplace_autoregression(
 208             model,
 209             self.batch_size,
 210             result,
 211             ar_mask,
 212             deterministic_synthesis,
 213             progress_bar_desc=None,
 214             device=self.device,
 215         )
 216
 217         self.save_image(
 218             result[:96],
 219             result_dir,
 220             f"world_prediction_{n_epoch:04d}_{model.id:02d}.png",
 221             logger,
 222         )
 223
 224         return main_test_accuracy
 225
 226     def store_new_quizzes(self, new_quizzes, for_train=True):
 227         if for_train:
 228             self.train_quizzes.append(new_quizzes)
 229         else:
 230             self.test_quizzes.append(new_quizzes)
 231
 232     def create_new_quizzes(
 233         self,
 234         n_epoch,
 235         result_dir,
 236         logger,
 237         nb,
 238         model,
 239         other_models,
 240     ):
 241         new_quizzes = torch.empty(
 242             nb, self.height * self.width * 2 + 1, device=self.device, dtype=torch.int64
 243         )
 244         ar_mask = torch.full(new_quizzes.size(), 1, device=self.device)
 245
 246         masked_inplace_autoregression(
 247             model,
 248             self.batch_size,
 249             new_quizzes,
 250             ar_mask,
 251             deterministic_synthesis=False,
 252             progress_bar_desc="creating quizzes",
 253             device=self.device,
 254         )
 255
 256         ar_mask = self.make_ar_mask(new_quizzes)
 257
 258         nb_correct = 0
 259
 260         for m in other_models:
 261             result = new_quizzes.clone()
 262
 263             masked_inplace_autoregression(
 264                 m,
 265                 self.batch_size,
 266                 result,
 267                 ar_mask,
 268                 deterministic_synthesis=True,
 269                 progress_bar_desc="solving quizzes",
 270                 device=self.device,
 271             )
 272
 273             l = self.height * self.width
 274             direction = new_quizzes[:, l : l + 1]
 275             direction = world.token_forward * (
 276                 direction == world.token_backward
 277             ) + world.token_backward * (direction == world.token_forward)
 278             inverted_quizzes = torch.cat(
 279                 [new_quizzes[:, l + 1 :], direction, new_quizzes[:, :l]], dim=1
 280             )
 281
 282             inverted_result = inverted_quizzes.clone()
 283
 284             masked_inplace_autoregression(
 285                 m,
 286                 self.batch_size,
 287                 inverted_result,
 288                 ar_mask,
 289                 deterministic_synthesis=True,
 290                 progress_bar_desc="solving reversed quizzes",
 291                 device=self.device,
 292             )
 293
 294             nb_correct += (new_quizzes == result).long().min(dim=-1).values * (
 295                 inverted_quizzes == inverted_result
 296             ).long().min(dim=-1).values
 297
 298         return new_quizzes, nb_correct