def load_svhn_infer(split: str, batch_size: int,
n_workers: Union[str, int]) -> DataLoader:
if split not in {"train", "val", "test"}:
raise ValueError("Split must be 'train', 'val', or 'test'!")
if batch_size <= 0:
raise ValueError("Batch_size must be positive!")
if type(n_workers) == str and n_workers != "n_cores":
raise ValueError("If n_workers is a string, it must be 'n_cores'!")
if type(n_workers) == int and n_workers < 0:
raise ValueError("If n_workers is an int, it must be non-negative!")
transform = transforms.ToTensor()
if split == "train":
dataset = ConcatDataset([
torchvision.datasets.SVHN(DATA_DIRPATH,
split="train",
transform=transform,
download=True),
torchvision.datasets.SVHN(DATA_DIRPATH,
split="extra",
transform=transform,
download=True)
])
indices = read_lines(os.path.join(SPLIT_DIRPATH, "train.txt"), int)
dataset = Subset(dataset, indices)
elif split == "val":
dataset = ConcatDataset([
torchvision.datasets.SVHN(DATA_DIRPATH,
split="train",
transform=transform,
download=True),
torchvision.datasets.SVHN(DATA_DIRPATH,
split="extra",
transform=transform,
download=True)
])
indices = read_lines(os.path.join(SPLIT_DIRPATH, "val.txt"), int)
dataset = Subset(dataset, indices)
else:
dataset = torchvision.datasets.SVHN(DATA_DIRPATH,
split="test",
transform=transform,
download=True)
if n_workers == "n_cores":
n_workers = multiprocessing.cpu_count()
return DataLoader(dataset,
batch_size=batch_size,
shuffle=False,
num_workers=n_workers)
def load_cifar10_train(batch_size: int, n_workers: Union[str,
int]) -> DataLoader:
if batch_size <= 0:
raise ValueError("Batch_size must be positive!")
if type(n_workers) == str and n_workers != "n_cores":
raise ValueError("If n_workers is a string, it must be 'n_cores'!")
if type(n_workers) == int and n_workers < 0:
raise ValueError("If n_workers is an int, it must be non-negative!")
transform = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.RandomCrop(32, 4),
transforms.ToTensor()
])
dataset = torchvision.datasets.CIFAR10(DATA_DIRPATH,
train=True,
transform=transform,
download=True)
train_indices = read_lines(os.path.join(SPLIT_DIRPATH, "train.txt"), int)
train_dataset = Subset(dataset, train_indices)
if n_workers == "n_cores":
n_workers = multiprocessing.cpu_count()
return DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=n_workers)
def setUp(self) -> None:
transform = torchvision.transforms.ToTensor()
dataset = ConcatDataset([
torchvision.datasets.SVHN(DATA_DIRPATH, split="train", transform=transform, download=True),
torchvision.datasets.SVHN(DATA_DIRPATH, split="extra", transform=transform, download=True)
])
train_indices = read_lines(os.path.join(SPLIT_DIRPATH, "train.txt"), int)
self.train_dataset = Subset(dataset, train_indices)
self.images = torch.stack([image for image, label in self.train_dataset], dim=0)
def setUp(self) -> None:
transform = torchvision.transforms.ToTensor()
dataset = torchvision.datasets.CIFAR10(DATA_DIRPATH,
train=True,
transform=transform,
download=True)
train_indices = read_lines(os.path.join(SPLIT_DIRPATH, "train.txt"),
int)
self.train_dataset = Subset(dataset, train_indices)
self.images = torch.stack(
[image for image, label in self.train_dataset], dim=0)