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         ###############################################################
 242         # Generate quizzes with model
 243
 244         quizzes = torch.empty(
 245             nb, self.height * self.width * 2 + 1, device=self.device, dtype=torch.int64
 246         )
 247         ar_mask = torch.full(quizzes.size(), 1, device=self.device)
 248
 249         masked_inplace_autoregression(
 250             model,
 251             self.batch_size,
 252             quizzes,
 253             ar_mask,
 254             deterministic_synthesis=False,
 255             progress_bar_desc="creating quizzes",
 256             device=self.device,
 257         )
 258
 259         ###############################################################
 260         # Create the reverse quizzes
 261
 262         l = self.height * self.width
 263         direction = quizzes[:, l : l + 1]
 264         direction = world.token_forward * (
 265             direction == world.token_backward
 266         ) + world.token_backward * (direction == world.token_forward)
 267         reverse_quizzes = torch.cat(
 268             [quizzes[:, l + 1 :], direction, quizzes[:, :l]], dim=1
 269         )
 270
 271         ar_mask = self.make_ar_mask(quizzes)
 272
 273         ###############################################################
 274         # Check how many of the other models can solve them in both
 275         # directions
 276
 277         nb_correct = 0
 278
 279         for m in other_models:
 280             result = quizzes.clone()
 281
 282             masked_inplace_autoregression(
 283                 m,
 284                 self.batch_size,
 285                 result,
 286                 ar_mask,
 287                 deterministic_synthesis=True,
 288                 progress_bar_desc="solving quizzes",
 289                 device=self.device,
 290             )
 291
 292             correct = (quizzes == result).long().min(dim=-1).values
 293
 294             reverse_result = reverse_quizzes.clone()
 295
 296             masked_inplace_autoregression(
 297                 m,
 298                 self.batch_size,
 299                 reverse_result,
 300                 ar_mask,
 301                 deterministic_synthesis=True,
 302                 progress_bar_desc="solving reversed quizzes",
 303                 device=self.device,
 304             )
 305
 306             reverse_correct = (
 307                 (reverse_quizzes == reverse_result).long().min(dim=-1).values
 308             )
 309
 310             nb_correct += correct * reverse_correct
 311
 312         return quizzes, nb_correct