X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=tasks.py;h=5583fc89827d82be551db14ac9cb601f670b4233;hb=3dea181a5903a0e577e4830c66405b40f2a2df1d;hp=75781ab7f9f3fab058ceae1ffe5d4653fb108656;hpb=c03e968adc7bf73df07a0fad89a835b98f4e76df;p=picoclvr.git diff --git a/tasks.py b/tasks.py index 75781ab..5583fc8 100755 --- a/tasks.py +++ b/tasks.py @@ -20,6 +20,8 @@ def masked_inplace_autoregression( progress_bar_desc="autoregression", device=torch.device("cpu"), ): + assert input.size() == ar_mask.size() + batches = zip(input.split(batch_size), ar_mask.split(batch_size)) if progress_bar_desc is not None: @@ -27,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(): @@ -58,6 +60,83 @@ 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 @@ -106,6 +185,8 @@ class PicoCLVR(Task): pruner_train=None, pruner_eval=None, ): + super().__init__() + def generate_descr(nb, cache_suffix, pruner): return picoclvr.generate( nb, @@ -294,6 +375,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 @@ -364,6 +447,8 @@ class Maze(Task): nb_walls, device=torch.device("cpu"), ): + super().__init__() + self.batch_size = batch_size self.height = height self.width = width @@ -535,6 +620,8 @@ class Snake(Task): prompt_length, device=torch.device("cpu"), ): + super().__init__() + self.batch_size = batch_size self.height = height self.width = width @@ -590,8 +677,6 @@ class Snake(Task): ) result *= 1 - ar_mask - # snake.solver(result,ar_mask) - masked_inplace_autoregression( model, self.batch_size, @@ -605,19 +690,8 @@ class Snake(Task): nb_correct = ((result == input).long() * (prior_visits > 0) * ar_mask).sum() - # nb_total = result.size(0) - # nb_correct = ((result - input).abs().sum(1) == 0).sum() - return nb_total, nb_correct - # train_nb_total, train_nb_correct = compute_nb_correct( - # self.train_input, self.train_prior_visits - # ) - - # logger( - # f"accuracy_train nb_total {train_nb_total} nb_correct {train_nb_correct} accuracy {(100.0*train_nb_correct)/train_nb_total:.02f}%" - # ) - test_nb_total, test_nb_correct = compute_nb_correct( self.test_input[:1000], self.test_prior_visits[:1000] ) @@ -646,6 +720,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 @@ -793,6 +869,8 @@ class Expr(Task): batch_size, device=torch.device("cpu"), ): + super().__init__() + self.batch_size = batch_size self.device = device @@ -957,3 +1035,128 @@ class Expr(Task): ###################################################################### + +import world + + +class World(Task): + def __init__( + self, + nb_train_samples, + 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, + train_action_seq, + test_frames, + test_action_seq, + self.frame2seq, + self.seq2frame, + ) = world.create_data_and_processors( + nb_train_samples, + nb_test_samples, + mode="first_last", + nb_steps=30, + nb_epochs=vqae_nb_epochs, + logger=logger, + device=device, + device_storage=device_storage, + ) + + print(f"{train_action_seq.size()=}") + + 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"} + 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.to(self.device) + + def vocabulary_size(self): + return self.nb_codes + + def produce_results( + self, n_epoch, model, result_dir, logger, deterministic_synthesis + ): + k = torch.arange( + 2 * self.len_frame_seq + self.len_action_seq, device=self.device + )[None, :] + + 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( + model, + self.batch_size, + result, + ar_mask, + deterministic_synthesis, + device=self.device, + ) + + 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=12, + padding=1, + pad_value=0.0, + ) + logger(f"wrote {image_name}") + + +######################################################################