X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=tasks.py;h=9cd06ae054ae7e1adee634a9361adb8680d1356c;hb=0f580d4facb4b4b485d0a38d62d06c0639715b77;hp=96d062169467ec6b2976e5c5a39e0c9a78537b9f;hpb=8e23dd068df00df61c690ffa89ecc8cb9db4b32d;p=picoclvr.git diff --git a/tasks.py b/tasks.py index 96d0621..9cd06ae 100755 --- a/tasks.py +++ b/tasks.py @@ -29,7 +29,7 @@ def masked_inplace_autoregression( batches, dynamic_ncols=True, desc=progress_bar_desc, - #total=input.size(0) // batch_size, + # total=input.size(0) // batch_size, ) with torch.autograd.no_grad(): @@ -60,6 +60,84 @@ class Task: pass +###################################################################### + + +class Problem: + def generate(nb): + pass + + def perf(seq, logger): + pass + + +class ProblemByheart(Problem): + def __init__(self): + nb_seq, len_prompt, len_result = 100, 5, 5 + self.seq = torch.randint(10, (nb_seq, len_prompt + 1 + len_result)) + self.seq[:, len_prompt] = -1 + + def generate_sequences(self, nb): + return self.seq[torch.randint(self.seq.size(0), (nb,))] + + +class SandBox(Task): + def __init__( + self, + nb_train_samples, + nb_test_samples, + batch_size, + logger=None, + device=torch.device("cpu"), + ): + super().__init__() + + self.batch_size = batch_size + + problems = [ProblemByheart()] + nb_common_codes = 100 + + def generate_sequences(nb_samples): + problem_indexes = torch.randint(len(problems), (nb_samples,)) + nb_samples_per_problem = torch.one_hot(problem_indexes).sum(0) + print(f"{nb_samples_per_problem}") + all_seq = [] + for nb, p in zip(nb_samples_per_problem, problems): + all_seq.append(p.generate_sequences(nb_samples_per_problem[nb])) + return all_seq + + train_seq = generate_sequences(nb_train_samples) + test_seq = generate_sequences(nb_test_samples) + + for strain, stest in zip(train_seq, test_seq): + s = torch.cat((strain, stest), 0) + + self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1 + + def batches(self, split="train", nb_to_use=-1, desc=None): + assert split in {"train", "test"} + input = self.train_input if split == "train" else self.test_input + if nb_to_use > 0: + input = input[:nb_to_use] + if desc is None: + desc = f"epoch-{split}" + for batch in tqdm.tqdm( + input.split(self.batch_size), dynamic_ncols=True, desc=desc + ): + yield batch + + def vocabulary_size(self): + return self.nb_codes + + def produce_results( + self, n_epoch, model, result_dir, logger, deterministic_synthesis + ): + # logger( + # 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}%" + # ) + pass + + ###################################################################### import picoclvr @@ -108,6 +186,8 @@ class PicoCLVR(Task): pruner_train=None, pruner_eval=None, ): + super().__init__() + def generate_descr(nb, cache_suffix, pruner): return picoclvr.generate( nb, @@ -296,6 +376,8 @@ class MNIST(Task): def __init__( self, nb_train_samples, nb_test_samples, batch_size, device=torch.device("cpu") ): + super().__init__() + self.nb_train_samples = (nb_train_samples,) self.nb_test_samples = (nb_test_samples,) self.batch_size = batch_size @@ -366,6 +448,8 @@ class Maze(Task): nb_walls, device=torch.device("cpu"), ): + super().__init__() + self.batch_size = batch_size self.height = height self.width = width @@ -537,6 +621,8 @@ class Snake(Task): prompt_length, device=torch.device("cpu"), ): + super().__init__() + self.batch_size = batch_size self.height = height self.width = width @@ -635,6 +721,8 @@ class Stack(Task): fraction_values_for_train=None, device=torch.device("cpu"), ): + super().__init__() + self.batch_size = batch_size self.nb_steps = nb_steps self.nb_stacks = nb_stacks @@ -782,6 +870,8 @@ class Expr(Task): batch_size, device=torch.device("cpu"), ): + super().__init__() + self.batch_size = batch_size self.device = device @@ -957,16 +1047,20 @@ class World(Task): nb_test_samples, batch_size, vqae_nb_epochs, + logger=None, device=torch.device("cpu"), + device_storage=torch.device("cpu"), ): + super().__init__() + self.batch_size = batch_size self.device = device ( train_frames, - self.train_actions, + train_action_seq, test_frames, - self.test_actions, + test_action_seq, self.frame2seq, self.seq2frame, ) = world.create_data_and_processors( @@ -975,15 +1069,35 @@ class World(Task): mode="first_last", nb_steps=30, nb_epochs=vqae_nb_epochs, + logger=logger, device=device, + device_storage=device_storage, ) - self.train_input = self.frame2seq(train_frames) - self.train_input = self.train_input.reshape(self.train_input.size(0) // 2, -1) - self.test_input = self.frame2seq(test_frames) - self.test_input = self.test_input.reshape(self.test_input.size(0) // 2, -1) + print(f"{train_action_seq.size()=}") - self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1 + train_frame_seq = self.frame2seq(train_frames).to(device_storage) + test_frame_seq = self.frame2seq(test_frames).to(device_storage) + + nb_frame_codes = max(train_frame_seq.max(), test_frame_seq.max()) + 1 + nb_action_codes = max(train_action_seq.max(), test_action_seq.max()) + 1 + + self.len_frame_seq = train_frame_seq.size(1) + self.len_action_seq = train_action_seq.size(1) + self.nb_codes = nb_frame_codes + nb_action_codes + + train_frame_seq = train_frame_seq.reshape(train_frame_seq.size(0) // 2, 2, -1) + print(f"{train_action_seq.device=} {nb_frame_codes.device=}") + train_action_seq += nb_frame_codes + self.train_input = torch.cat( + (train_frame_seq[:, 0, :], train_action_seq, train_frame_seq[:, 1, :]), 1 + ) + + test_frame_seq = test_frame_seq.reshape(test_frame_seq.size(0) // 2, 2, -1) + test_action_seq += nb_frame_codes + self.test_input = torch.cat( + (test_frame_seq[:, 0, :], test_action_seq, test_frame_seq[:, 1, :]), 1 + ) def batches(self, split="train", nb_to_use=-1, desc=None): assert split in {"train", "test"} @@ -995,7 +1109,7 @@ class World(Task): for batch in tqdm.tqdm( input.split(self.batch_size), dynamic_ncols=True, desc=desc ): - yield batch + yield batch.to(self.device) def vocabulary_size(self): return self.nb_codes @@ -1003,11 +1117,16 @@ class World(Task): def produce_results( self, n_epoch, model, result_dir, logger, deterministic_synthesis ): - l = self.train_input.size(1) - k = torch.arange(l, device=self.device)[None, :] - result = self.test_input[:64].clone() + k = torch.arange( + 2 * self.len_frame_seq + self.len_action_seq, device=self.device + )[None, :] - ar_mask = (k >= l // 2).long().expand_as(result) + input = self.test_input[:64].to(self.device) + result = input.clone() + + ar_mask = ( + (k >= self.len_frame_seq + self.len_action_seq).long().expand_as(result) + ) result *= 1 - ar_mask masked_inplace_autoregression( @@ -1019,14 +1138,22 @@ class World(Task): device=self.device, ) - result = result.reshape(result.size(0) * 2, -1) + seq_start = input[:, : self.len_frame_seq] + seq_end = input[:, self.len_frame_seq + self.len_action_seq :] + seq_predicted = result[:, self.len_frame_seq + self.len_action_seq :] + + result = torch.cat( + (seq_start[:, None, :], seq_end[:, None, :], seq_predicted[:, None, :]), 1 + ) + result = result.reshape(-1, result.size(-1)) + print(f"{result.size()=}") frames = self.seq2frame(result) image_name = os.path.join(result_dir, f"world_result_{n_epoch:04d}.png") torchvision.utils.save_image( frames.float() / (world.Box.nb_rgb_levels - 1), image_name, - nrow=8, + nrow=12, padding=1, pad_value=0.0, )