3 # Any copyright is dedicated to the Public Domain.
4 # https://creativecommons.org/publicdomain/zero/1.0/
6 # Written by Francois Fleuret <francois@fleuret.org>
8 import math, os, tqdm, warnings
10 import torch, torchvision
13 from torch.nn import functional as F
15 from mygpt import BracketedSequence
17 # from graph import save_attention_image
18 save_attention_image = None
20 ######################################################################
23 def masked_inplace_autoregression(
28 deterministic_synthesis,
29 forbidden_tokens=None,
31 progress_bar_desc="autoregression",
32 device=torch.device("cpu"),
34 assert input.size() == ar_mask.size()
36 batches = zip(input.split(batch_size), ar_mask.split(batch_size))
38 if progress_bar_desc is not None:
42 desc=progress_bar_desc,
43 total=(input.size(0) + batch_size - 1) // batch_size,
46 with torch.autograd.no_grad():
50 for input, ar_mask in batches:
51 model.masked_inplace_autoregression(
54 deterministic_synthesis,
62 ######################################################################
66 def batches(self, split="train", nb_to_use=-1, desc=None):
69 def vocabulary_size(self):
73 self, n_epoch, model, result_dir, logger, deterministic_synthesis
78 ######################################################################
84 def save_image(self, input, result_dir, filename, logger):
85 img = world.sample2img(input.to("cpu"), self.height, self.width)
86 image_name = os.path.join(result_dir, filename)
87 torchvision.utils.save_image(img.float() / 255.0, image_name, nrow=8, padding=2)
88 logger(f"wrote {image_name}")
90 def make_ar_mask(self, input):
91 b = torch.arange(input.size(1), device=input.device) > input.size(1) // 2
92 return b.long()[None, :].expand_as(input)
101 device=torch.device("cpu"),
105 self.batch_size = batch_size
110 self.train_input = world.generate(
111 nb_train_samples, height=self.height, width=self.width
114 self.test_input = world.generate(
115 nb_test_samples, height=self.height, width=self.width
118 self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1
120 self.train_quizzes = []
121 self.test_quizzes = []
123 if result_dir is not None:
125 self.train_input[:96], result_dir, f"world_train.png", logger
128 def batches(self, split="train", desc=None):
129 assert split in {"train", "test"}
131 input = self.train_input
132 quizzes = self.train_quizzes
134 input = self.test_input
135 quizzes = self.test_quizzes
138 quizzes = torch.cat(quizzes, dim=0)
139 if quizzes.size(0) > input.size(0) // 2:
140 i = torch.randperm(input.size(0))[: input.size(0) // 2]
143 i = torch.randperm(input.size(0))[: input.size(0) - quizzes.size(0)]
146 self.nb_batch_samples_world = input.size(0)
147 self.nb_batch_samples_quizzes = quizzes.size(0)
149 input = torch.cat([input, quizzes], dim=0)
151 self.nb_batch_samples_world = input.size(0)
152 self.nb_batch_samples_quizzes = 0
155 desc = f"epoch-{split}"
156 for batch in tqdm.tqdm(
157 input.split(self.batch_size), dynamic_ncols=True, desc=desc
161 def vocabulary_size(self):
165 self, n_epoch, model, result_dir, logger, deterministic_synthesis, nmax=1000
167 def compute_accuracy(input, logger=None):
169 ar_mask = self.make_ar_mask(input)
170 result = input.clone() * (1 - ar_mask)
172 masked_inplace_autoregression(
177 deterministic_synthesis,
178 progress_bar_desc=None,
182 nb_total, nb_correct = (
184 (input == result).long().min(dim=1).values.sum(),
187 return nb_total, nb_correct
189 train_nb_total, train_nb_correct = compute_accuracy(self.train_input)
192 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}%"
195 test_nb_total, test_nb_correct = compute_accuracy(self.test_input, logger)
198 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}%"
201 main_test_accuracy = test_nb_correct / test_nb_total
202 logger(f"main_test_accuracy {n_epoch} {main_test_accuracy}")
204 ##############################
206 input = self.test_input[:96]
207 ar_mask = self.make_ar_mask(input)
208 result = input.clone() * (1 - ar_mask)
210 masked_inplace_autoregression(
215 deterministic_synthesis,
216 progress_bar_desc=None,
223 f"world_result_{n_epoch:04d}_{model.id:02d}.png",
227 return main_test_accuracy
229 def store_new_quizzes(self, new_quizzes, for_train=True):
231 self.train_quizzes.append(new_quizzes)
233 self.test_quizzes.append(new_quizzes)
235 def create_new_quizzes(
244 new_quizzes = torch.empty(
245 nb, self.height * self.width * 2 + 1, device=self.device, dtype=torch.int64
247 ar_mask = torch.full(new_quizzes.size(), 1, device=self.device)
249 masked_inplace_autoregression(
254 deterministic_synthesis=False,
255 progress_bar_desc="new quizzes",
259 ar_mask = self.make_ar_mask(new_quizzes)
263 for m in other_models:
264 result = new_quizzes.clone()
266 masked_inplace_autoregression(
271 deterministic_synthesis=True,
272 progress_bar_desc="solving quizzes",
276 nb_correct += (new_quizzes == result).long().min(dim=-1).values
278 return new_quizzes, nb_correct