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=6, padding=4)
  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_seq(
 108             nb_train_samples, height=self.height, width=self.width
 109         ).to(device)
 110
 111         self.test_input = world.generate_seq(
 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[:72], 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[:72],
 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 renew_samples(self, nb, for_train=True):
 227         input = self.train_input if for_train else self.test_input
 228         nb = min(nb, input.size(0))
 229         input[:-nb] = input[nb:].clone()
 230         input[-nb:] = world.generate_seq(nb, height=self.height, width=self.width).to(
 231             self.device
 232         )
 233
 234     def store_new_quizzes(self, new_quizzes, for_train=True):
 235         if for_train:
 236             self.train_quizzes.append(new_quizzes)
 237         else:
 238             self.test_quizzes.append(new_quizzes)
 239
 240     def create_new_quizzes(
 241         self,
 242         n_epoch,
 243         result_dir,
 244         logger,
 245         nb,
 246         model,
 247         other_models,
 248     ):
 249         ###############################################################
 250         # Generate quizzes with model
 251
 252         quizzes = torch.empty(
 253             nb, self.height * self.width * 2 + 1, device=self.device, dtype=torch.int64
 254         )
 255         ar_mask = torch.full(quizzes.size(), 1, device=self.device)
 256
 257         masked_inplace_autoregression(
 258             model,
 259             self.batch_size,
 260             quizzes,
 261             ar_mask,
 262             deterministic_synthesis=False,
 263             progress_bar_desc="creating quizzes",
 264             device=self.device,
 265         )
 266
 267         ###############################################################
 268         # Create the reverse quizzes
 269
 270         l = self.height * self.width
 271         direction = quizzes[:, l : l + 1]
 272         direction = world.token_forward * (
 273             direction == world.token_backward
 274         ) + world.token_backward * (direction == world.token_forward)
 275         reverse_quizzes = torch.cat(
 276             [quizzes[:, l + 1 :], direction, quizzes[:, :l]], dim=1
 277         )
 278
 279         ar_mask = self.make_ar_mask(quizzes)
 280
 281         ###############################################################
 282         # Check how many of the other models can solve them in both
 283         # directions
 284
 285         nb_correct = []
 286
 287         for m in other_models:
 288             result = quizzes.clone()
 289
 290             masked_inplace_autoregression(
 291                 m,
 292                 self.batch_size,
 293                 result,
 294                 ar_mask,
 295                 deterministic_synthesis=True,
 296                 progress_bar_desc="solving quizzes",
 297                 device=self.device,
 298             )
 299
 300             correct = (quizzes == result).long().min(dim=-1).values
 301
 302             reverse_result = reverse_quizzes.clone()
 303
 304             masked_inplace_autoregression(
 305                 m,
 306                 self.batch_size,
 307                 reverse_result,
 308                 ar_mask,
 309                 deterministic_synthesis=True,
 310                 progress_bar_desc="solving reversed quizzes",
 311                 device=self.device,
 312             )
 313
 314             reverse_correct = (
 315                 (reverse_quizzes == reverse_result).long().min(dim=-1).values
 316             )
 317
 318             nb_correct.append((correct * reverse_correct)[None, :])
 319
 320         nb_correct = torch.cat(nb_correct, dim=0)
 321
 322         filename = os.path.join(result_dir, "correct_{n_epoch:04d}.dat")
 323         with open(filename, "w") as f:
 324             for k in nb_correct:
 325                 f.write(f"{k}\n")
 326
 327         return quizzes, nb_correct.sum(dim=0)