From: François Fleuret Date: Mon, 19 Jun 2023 16:51:17 +0000 (+0200) Subject: Update. X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=commitdiff_plain;h=f04e158f0679195821a7288bbfe08f775b894096;p=picoclvr.git Update. --- diff --git a/main.py b/main.py index ae42544..24c21f5 100755 --- a/main.py +++ b/main.py @@ -451,13 +451,53 @@ class TaskPicoCLVR(Task): 0, ) - image_name = os.path.join(args.result_dir, f"result_{n_epoch:04d}.png") + image_name = os.path.join(args.result_dir, f"picoclvr_result_{n_epoch:04d}.png") torchvision.utils.save_image( img / 255.0, image_name, nrow=nb_per_primer, padding=1, pad_value=1.0 ) log_string(f"wrote {image_name}") +###################################################################### + + +class TaskMNIST(Task): + def __init__(self, batch_size, device=torch.device("cpu")): + self.device = device + self.batch_size = batch_size + + def batches(self, split="train"): + assert split in {"train", "test"} + data_set = torchvision.datasets.MNIST( + root="./data", train=(split == "train"), download=True + ) + data_input = data_set.data.view(-1, 28 * 28).long() + if args.nb_train_samples is not None: + data_input = data_input[: args.nb_train_samples] + for batch in tqdm.tqdm( + data_input.split(self.batch_size), desc=f"epoch-{split}" + ): + yield batch + + def vocabulary_size(self): + return 256 + + def produce_results(self, n_epoch, model): + results = torch.empty(64, 28 * 28, device=self.device, dtype=torch.int64) + ar_mask = torch.full_like(results, 1) + masked_inplace_autoregression( + model, self.batch_size, results, ar_mask, device=self.device + ) + image_name = os.path.join(args.result_dir, f"result_mnist_{n_epoch:04d}.png") + torchvision.utils.save_image( + 1 - results.reshape(-1, 1, 28, 28) / 255.0, + image_name, + nrow=16, + pad_value=0.8, + ) + log_string(f"wrote {image_name}") + + ###################################################################### import maze @@ -589,9 +629,10 @@ class TaskMaze(Task): mazes, paths = self.seq2map(input) _, predicted_paths = self.seq2map(result) - filename = f"result_{n_epoch:04d}.png" + + filename = os.path.join(args.result_dir, f"result_{n_epoch:04d}.png") maze.save_image( - os.path.join(args.result_dir, filename), + filename, mazes=mazes, target_paths=paths, predicted_paths=predicted_paths, @@ -636,6 +677,12 @@ if args.task == "picoclvr": pruner_eval=picoclvr_pruner_eval, ) +elif args.task == "mnist": + task = TaskMNIST( + batch_size=args.batch_size, + device=device, + ) + elif args.task == "maze": task = TaskMaze( nb_train_samples=args.nb_train_samples,