+######################################################################
+
+
+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}")
+
+