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         # print()
 116         # for a in world.seq2str(self.train_input):
 117         # print(a)
 118         # for a in world.seq2str(self.test_input):
 119         # print(a)
 120         # exit(0)
 121
 122         self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1
 123
 124         self.train_quizzes = []
 125         self.test_quizzes = []
 126
 127         if result_dir is not None:
 128             self.save_image(
 129                 self.train_input[:96], result_dir, f"world_train.png", logger
 130             )
 131
 132     def batches(self, split="train", desc=None):
 133         assert split in {"train", "test"}
 134         if split == "train":
 135             input = self.train_input
 136             quizzes = self.train_quizzes
 137         else:
 138             input = self.test_input
 139             quizzes = self.test_quizzes
 140
 141         if len(quizzes) > 0:
 142             quizzes = torch.cat(quizzes, dim=0)
 143             if quizzes.size(0) > input.size(0) // 2:
 144                 i = torch.randperm(input.size(0))[: input.size(0) // 2]
 145                 quizzes = quizzes[i]
 146
 147             i = torch.randperm(input.size(0))[: input.size(0) - quizzes.size(0)]
 148             input = input[i]
 149
 150             self.nb_batch_samples_world = input.size(0)
 151             self.nb_batch_samples_quizzes = quizzes.size(0)
 152
 153             input = torch.cat([input, quizzes], dim=0)
 154         else:
 155             self.nb_batch_samples_world = input.size(0)
 156             self.nb_batch_samples_quizzes = 0
 157
 158         if desc is None:
 159             desc = f"epoch-{split}"
 160         for batch in tqdm.tqdm(
 161             input.split(self.batch_size), dynamic_ncols=True, desc=desc
 162         ):
 163             yield batch
 164
 165     def vocabulary_size(self):
 166         return self.nb_codes
 167
 168     def produce_results(
 169         self, n_epoch, model, result_dir, logger, deterministic_synthesis, nmax=1000
 170     ):
 171         def compute_accuracy(input, logger=None):
 172             input = input[:nmax]
 173             ar_mask = self.make_ar_mask(input)
 174             result = input.clone() * (1 - ar_mask)
 175
 176             masked_inplace_autoregression(
 177                 model,
 178                 self.batch_size,
 179                 result,
 180                 ar_mask,
 181                 deterministic_synthesis,
 182                 progress_bar_desc=None,
 183                 device=self.device,
 184             )
 185
 186             nb_total, nb_correct = (
 187                 input.size(0),
 188                 (input == result).long().min(dim=1).values.sum(),
 189             )
 190
 191             return nb_total, nb_correct
 192
 193         train_nb_total, train_nb_correct = compute_accuracy(self.train_input)
 194
 195         logger(
 196             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}%"
 197         )
 198
 199         test_nb_total, test_nb_correct = compute_accuracy(self.test_input, logger)
 200
 201         logger(
 202             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}%"
 203         )
 204
 205         main_test_accuracy = test_nb_correct / test_nb_total
 206         logger(f"main_test_accuracy {n_epoch} {main_test_accuracy}")
 207
 208         ##############################
 209
 210         input = self.test_input[:96]
 211         ar_mask = self.make_ar_mask(input)
 212         result = input.clone() * (1 - ar_mask)
 213
 214         masked_inplace_autoregression(
 215             model,
 216             self.batch_size,
 217             result,
 218             ar_mask,
 219             deterministic_synthesis,
 220             progress_bar_desc=None,
 221             device=self.device,
 222         )
 223
 224         self.save_image(
 225             result[:96],
 226             result_dir,
 227             f"world_prediction_{n_epoch:04d}_{model.id:02d}.png",
 228             logger,
 229         )
 230
 231         return main_test_accuracy
 232
 233     def store_new_quizzes(self, new_quizzes, for_train=True):
 234         if for_train:
 235             self.train_quizzes.append(new_quizzes)
 236         else:
 237             self.test_quizzes.append(new_quizzes)
 238
 239     def create_new_quizzes(
 240         self,
 241         n_epoch,
 242         result_dir,
 243         logger,
 244         nb,
 245         model,
 246         other_models,
 247     ):
 248         ###############################################################
 249         # Generate quizzes with model
 250
 251         quizzes = torch.empty(
 252             nb, self.height * self.width * 2 + 1, device=self.device, dtype=torch.int64
 253         )
 254         ar_mask = torch.full(quizzes.size(), 1, device=self.device)
 255
 256         masked_inplace_autoregression(
 257             model,
 258             self.batch_size,
 259             quizzes,
 260             ar_mask,
 261             deterministic_synthesis=False,
 262             progress_bar_desc="creating quizzes",
 263             device=self.device,
 264         )
 265
 266         ###############################################################
 267         # Create the reverse quizzes
 268
 269         l = self.height * self.width
 270         direction = quizzes[:, l : l + 1]
 271         direction = world.token_forward * (
 272             direction == world.token_backward
 273         ) + world.token_backward * (direction == world.token_forward)
 274         reverse_quizzes = torch.cat(
 275             [quizzes[:, l + 1 :], direction, quizzes[:, :l]], dim=1
 276         )
 277
 278         ar_mask = self.make_ar_mask(quizzes)
 279
 280         ###############################################################
 281         # Check how many of the other models can solve them in both
 282         # directions
 283
 284         nb_correct = []
 285
 286         for m in other_models:
 287             result = quizzes.clone()
 288
 289             masked_inplace_autoregression(
 290                 m,
 291                 self.batch_size,
 292                 result,
 293                 ar_mask,
 294                 deterministic_synthesis=True,
 295                 progress_bar_desc="solving quizzes",
 296                 device=self.device,
 297             )
 298
 299             correct = (quizzes == result).long().min(dim=-1).values
 300
 301             reverse_result = reverse_quizzes.clone()
 302
 303             masked_inplace_autoregression(
 304                 m,
 305                 self.batch_size,
 306                 reverse_result,
 307                 ar_mask,
 308                 deterministic_synthesis=True,
 309                 progress_bar_desc="solving reversed quizzes",
 310                 device=self.device,
 311             )
 312
 313             reverse_correct = (
 314                 (reverse_quizzes == reverse_result).long().min(dim=-1).values
 315             )
 316
 317             nb_correct.append((correct * reverse_correct)[None, :])
 318
 319         nb_correct = torch.cat(nb_correct, dim=0)
 320
 321         filename = os.path.join(result_dir, "correct_{n_epoch:04d}.dat")
 322         with open(filename, "w") as f:
 323             for k in nb_correct:
 324                 f.write(f"{k}\n")
 325
 326         return quizzes, nb_correct.sum(dim=0)