mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-18 07:41:03 +00:00
Brandon Thomas
135 lines
4.4 KiB
Python
135 lines
4.4 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from dataclasses import dataclass
|
|
from os.path import expanduser
|
|
from typing import Optional
|
|
|
|
from torchvision.datasets import MNIST
|
|
from torchvision.transforms import Compose, Resize, ToTensor
|
|
|
|
from nemo.core.classes import Dataset
|
|
from nemo.core.config import Config
|
|
from nemo.core.neural_types.axes import AxisKind, AxisType
|
|
from nemo.core.neural_types.elements import ClassificationTarget, Index, NormalizedImageValue, StringLabel
|
|
from nemo.core.neural_types.neural_type import NeuralType
|
|
from nemo.utils.decorators import add_port_docs, experimental
|
|
|
|
|
|
@dataclass
|
|
class MNISTDatasetConfig(Config):
|
|
"""
|
|
Structured config for MNISTDataset class.
|
|
|
|
Args:
|
|
height: image height (DEFAULT: 28)
|
|
width: image width (DEFAULT: 28)
|
|
data_folder: path to the folder with data, can be relative to user (DEFAULT: "~/data/mnist")
|
|
train: use train or test splits (DEFAULT: True)
|
|
name: Name of the module (DEFAULT: None)
|
|
"""
|
|
|
|
height: int = 28
|
|
width: int = 28
|
|
data_folder: str = "~/data/mnist"
|
|
train: bool = True
|
|
download: bool = True
|
|
|
|
|
|
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
|