class MNISTDataset(Dataset):
def __init__(self,
root_dir,
transform=transforms.Compose([transforms.ToTensor()])):
super().__init__()
self.root_dir = root_dir
self.transform = transform
self.dataset = MNIST(root_dir, transform=transform)
self.num_classes = len(self.dataset.classes)
def __getitem__(self, index):
x, y = self.dataset.__getitem__(index)
x = torch.cat((x, x, x), dim=0)
y_onehot = [0.] * self.num_classes
y_onehot[y] = 1.
return {"x": x, "y_onehot": np.asarray(y_onehot, dtype=np.float32)}
def __len__(self):
return len(self.dataset)
class MNISTDataset(Dataset):
"""
A "thin wrapper" around the torchvision's MNIST dataset.
"""
def __init__(self, cfg: MNISTDatasetConfig = MNISTDatasetConfig()):
"""
Initializes the MNIST dataset.
Args:
cfg: Configuration object of type MNISTDatasetConfig.
"""
# Call the base class constructor of Dataset.
Dataset.__init__(self) # , name=name)
# Store height and width.
self._height = cfg.height
self._width = cfg.width
# Create transformations: up-scale and transform to tensors.
mnist_transforms = Compose(
[Resize((self._height, self._width)),
ToTensor()])
# Get absolute path.
abs_data_folder = expanduser(cfg.data_folder)
# Create the MNIST dataset object.
self._dataset = MNIST(root=abs_data_folder,
train=cfg.train,
download=cfg.download,
transform=mnist_transforms)
# Class names.
labels = 'Zero One Two Three Four Five Six Seven Eight Nine'.split(' ')
word_to_ix = {labels[i]: i for i in range(10)}
# Reverse mapping.
self._ix_to_word = {value: key for (key, value) in word_to_ix.items()}
@property
def output_types(self):
"""
Creates definitions of output ports.
"""
return {
"indices":
NeuralType(tuple('B'), elements_type=Index()),
"images":
NeuralType(
axes=(
AxisType(kind=AxisKind.Batch),
AxisType(kind=AxisKind.Channel, size=1),
AxisType(kind=AxisKind.Height, size=self._height),
AxisType(kind=AxisKind.Width, size=self._width),
),
elements_type=NormalizedImageValue(), # float, <0-1>
),
"targets":
NeuralType(
tuple('B'),
elements_type=ClassificationTarget()), # Target are ints!
"labels":
NeuralType(tuple('B'),
elements_type=StringLabel()), # Labels is string!
}
def __len__(self):
"""
Returns:
Length of the dataset.
"""
return len(self._dataset)
def __getitem__(self, index: int):
"""
Returns a single sample.
Args:
index: index of the sample to return.
"""
# Get image and target.
img, target = self._dataset.__getitem__(index)
# Return sample.
return index, img, target, self._ix_to_word[target]
@property
def ix_to_word(self):
"""
Returns:
Dictionary with mapping of target indices (int) to labels (class names as strings)
that can we used by other modules.
"""
return self._ix_to_word
