X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=tasks.py;h=183c3cfc0ff7faeae97a4ec2de2dd529ded3192b;hb=cb52b31a10ccf9b8df95114efb6a8039c1e006b6;hp=24c13fe79eddb4b8104bfe869c079bb4cbce80b9;hpb=5332c56acd44d7049f3fbb33a8643482e0c71f4d;p=picoclvr.git diff --git a/tasks.py b/tasks.py index 24c13fe..183c3cf 100755 --- a/tasks.py +++ b/tasks.py @@ -1426,7 +1426,7 @@ import grid class Grid(Task): # Make a tensor from a list of strings - def tensorize(self, descr): + def str2tensor(self, descr): token_descr = [s.strip().split(" ") for s in descr] l = max([len(s) for s in token_descr]) token_descr = [s + ["#"] * (l - len(s)) for s in token_descr] @@ -1434,7 +1434,7 @@ class Grid(Task): return torch.tensor(id_descr, device=self.device) # Make a list of strings from a tensor - def detensorize(self, x): + def tensor2str(self, x): return [" ".join([self.id2token[t.item()] for t in r]) for r in x] # trim all the tensors in the tuple z to remove as much token from @@ -1495,12 +1495,12 @@ class Grid(Task): self.token2id = dict([(t, n) for n, t in enumerate(tokens)]) self.id2token = dict([(n, t) for n, t in enumerate(tokens)]) self.t_nul = self.token2id["#"] - self.t_true = self.token2id[""] - self.t_false = self.token2id[""] + self.t_true = self.token2id["true"] + self.t_false = self.token2id["false"] # Tokenize the train and test sets - self.train_input = self.tensorize(self.train_descr) - self.test_input = self.tensorize(self.test_descr) + self.train_input = self.str2tensor(self.train_descr) + self.test_input = self.str2tensor(self.test_descr) def batches(self, split="train"): assert split in {"train", "test"} @@ -1519,9 +1519,11 @@ class Grid(Task): correct = self.test_input[:1000] result = correct.clone() ar_mask = torch.logical_or(result == self.t_true, result == self.t_false).long() - result *= 1 - ar_mask + result *= 1 - ar_mask # paraaaaanoiaaaaaaa - for e in self.detensorize(result[:10]): + logger(f"----------------------------------------------------------") + + for e in self.tensor2str(result[:10]): logger(f"test_before {e}") masked_inplace_autoregression( @@ -1533,8 +1535,12 @@ class Grid(Task): device=self.device, ) - for e in self.detensorize(result[:10]): - logger(f"test_after {e}") + logger(f"----------------------------------------------------------") + + for e in self.tensor2str(result[:10]): + logger(f"test_after {e}") + + logger(f"----------------------------------------------------------") nb_total = ar_mask.sum().item() nb_correct = ((correct == result).long() * ar_mask).sum().item() @@ -1544,125 +1550,3 @@ class Grid(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, - ) - - 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) - - 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)) - - 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}") - - -######################################################################