File size: 1,855 Bytes
3589275 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | import os
import PIL
import torch
import numpy as np
import torchvision
from torchvision import transforms
from torchvision.datasets import CIFAR10 as PyTorchCIFAR10
from torchvision.datasets import VisionDataset
cifar_classnames = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
class CIFAR10:
def __init__(self, preprocess,
location=os.path.expanduser('~/data'),
batch_size=128,
num_workers=16):
self.train_dataset = PyTorchCIFAR10(
root=location, download=True, train=True, transform=preprocess
)
self.train_loader = torch.utils.data.DataLoader(
self.train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers
)
self.test_dataset = PyTorchCIFAR10(
root=location, download=True, train=False, transform=preprocess
)
self.test_loader = torch.utils.data.DataLoader(
self.test_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers
)
self.classnames = self.test_dataset.classes
def convert(x):
if isinstance(x, np.ndarray):
return torchvision.transforms.functional.to_pil_image(x)
return x
class BasicVisionDataset(VisionDataset):
def __init__(self, images, targets, transform=None, target_transform=None):
if transform is not None:
transform.transforms.insert(0, convert)
super(BasicVisionDataset, self).__init__(root=None, transform=transform, target_transform=target_transform)
assert len(images) == len(targets)
self.images = images
self.targets = targets
def __getitem__(self, index):
return self.transform(self.images[index]), self.targets[index]
def __len__(self):
return len(self.targets)
|