X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=tasks.py;h=7a4abbeea23d9494d835cf6a040a36ab2eb53cc2;hb=07c40deef359a676eea9b063e8f5f8a2254cad13;hp=d787c59e7dce1e70ec6a5b2b386c955b4536b018;hpb=6681907dcc86bf6e159925814d419f522e0e3300;p=picoclvr.git diff --git a/tasks.py b/tasks.py index d787c59..7a4abbe 100755 --- a/tasks.py +++ b/tasks.py @@ -111,13 +111,19 @@ class SandBox(Task): self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1 # A bit of paranoia never hurts - assert ( - self.nb_codes <= max_nb_codes - and self.train_input.min() >= 0 - and self.test_input.min() >= 0 - and tuple(self.train_ar_mask.unique()) == (0, 1) - and tuple(self.test_ar_mask.unique()) == (0, 1) - ) + assert self.nb_codes <= max_nb_codes + assert self.train_input.min() >= 0 + assert self.test_input.min() >= 0 + assert tuple(x.item() for x in self.train_ar_mask.unique()) in { + (0,), + (1,), + (0, 1), + } + assert tuple(x.item() for x in self.test_ar_mask.unique()) in { + (0,), + (1,), + (0, 1), + } def batches(self, split="train", nb_to_use=-1, desc=None): assert split in {"train", "test"} @@ -151,17 +157,24 @@ class SandBox(Task): device=self.device, ) + log_ground_truth = ar_mask.min() == 0 + if logger is not None: for sp, st in zip(result[:10], input[:10]): logger( f"test_sequences {n_epoch} prediction {self.problem.seq2str(sp)}" ) - logger( - f" {n_epoch} ground truth {self.problem.seq2str(st)}" - ) + if log_ground_truth: + logger( + f" {n_epoch} ground truth {self.problem.seq2str(st)}" + ) - nb_total = ar_mask.sum().item() - nb_correct = ((result == input).long() * ar_mask).sum().item() + nb_total, nb_correct = self.problem.compute_nb_correct( + input, ar_mask, result + ) + + # nb_total = ar_mask.sum().item() + # nb_correct = ((result == input).long() * ar_mask).sum().item() return nb_total, nb_correct @@ -1551,77 +1564,67 @@ class Grid(Task): ###################################################################### -import world +import qmlp + +class QMLP(Task): + ###################### -class World(Task): def __init__( self, nb_train_samples, nb_test_samples, batch_size, - vqae_nb_epochs, + result_dir, logger=None, device=torch.device("cpu"), - device_storage=torch.device("cpu"), ): super().__init__() - self.batch_size = batch_size self.device = device + self.batch_size = batch_size + self.nb_samples_per_mlp = 256 - ( - 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) + if logger is not None: + logger( + f"generating {nb_train_samples+nb_test_samples} samples (can take some time)" + ) - 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 + seq, q_test_set, test_error = qmlp.generate_sequence_and_test_set( + nb_mlps=nb_train_samples + nb_test_samples, + nb_samples=self.nb_samples_per_mlp, + device=self.device, + batch_size=64, + nb_epochs=250, + nb_mlps_per_batch=1024, + ) - 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 + self.train_input = seq[:nb_train_samples] + self.train_q_test_set = q_test_set[:nb_train_samples] + self.train_ref_test_errors = test_error[:nb_train_samples] + self.test_input = seq[nb_train_samples:] + self.test_q_test_set = q_test_set[nb_train_samples:] + self.test_ref_test_errors = test_error[nb_train_samples:] - train_frame_seq = train_frame_seq.reshape(train_frame_seq.size(0) // 2, 2, -1) + filename = os.path.join(result_dir, f"train_errors_ref.dat") + with open(filename, "w") as f: + for e in self.train_ref_test_errors: + f.write(f"{e}\n") - train_action_seq += nb_frame_codes - self.train_input = torch.cat( - (train_frame_seq[:, 0, :], train_action_seq, train_frame_seq[:, 1, :]), 1 - ) + filename = os.path.join(result_dir, f"test_errors_ref.dat") + with open(filename, "w") as f: + for e in self.test_ref_test_errors: + f.write(f"{e}\n") - 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 - ) + self.nb_codes = max(self.train_input.max(), self.test_input.max()) + 1 - def batches(self, split="train", nb_to_use=-1, desc=None): + def batches(self, split="train"): 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 + input.split(self.batch_size), dynamic_ncols=True, desc=f"epoch-{split}" ): - yield batch.to(self.device) + yield batch def vocabulary_size(self): return self.nb_codes @@ -1629,17 +1632,14 @@ class World(Task): 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() - + correct = self.test_input[:1000] + result = correct.clone() ar_mask = ( - (k >= self.len_frame_seq + self.len_action_seq).long().expand_as(result) - ) - result *= 1 - ar_mask + torch.arange(result.size(1), device=result.device) + > self.nb_samples_per_mlp * 3 + 1 + ).long()[None, :] + ar_mask = ar_mask.expand_as(result) + result *= 1 - ar_mask # paraaaaanoiaaaaaaa masked_inplace_autoregression( model, @@ -1650,25 +1650,14 @@ class World(Task): 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 :] + q_train_set = result[:, : self.nb_samples_per_mlp * 3] + q_params = result[:, self.nb_samples_per_mlp * 3 + 1 :] + error_test = qmlp.evaluate_q_params(q_params, self.test_q_test_set) - 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}") + filename = os.path.join(result_dir, f"test_errors_{n_epoch:04d}.dat") + with open(filename, "w") as f: + for e in error_test: + f.write(f"{e}\n") ######################################################################