Class GetDataLoaderPoison

mlh/data_preprocessing/data_loader.py:224–333

class GetDataLoaderPoison(object):
    def __init__(self, args):
        self.args = args
        self.data_path = args.data_path
        self.input_shape = args.input_shape

    def parse_dataset(self, dataset):

        if dataset in configs.SUPPORTED_IMAGE_DATASETS:
            _loader = getattr(datasets, dataset)
            if dataset != "EMNIST":
                train_dataset = _loader(root=self.data_path,
                                        train=True,
                                        transform=None,
                                        download=True)
                test_dataset = _loader(root=self.data_path,
                                       train=False,
                                       transform=None,
                                       download=True)
            else:
                train_dataset = _loader(root=self.data_path,
                                        train=True,
                                        split="byclass",
                                        transform=None,
                                        download=True)
                test_dataset = _loader(root=self.data_path,
                                       train=False,
                                       split="byclass",
                                       transform=None,
                                       download=True)

        else:
            raise ValueError("Dataset Not Supported: ", dataset)
        return train_dataset, test_dataset

    def get_data_transform(self, dataset, use_transform="simple"):
        transform_list = [transforms.Resize(
            (self.input_shape[0], self.input_shape[0])), ]

        if use_transform == "simple":
            transform_list += [transforms.RandomCrop(
                32, padding=4), transforms.RandomHorizontalFlip(), ]

            print("add simple data augmentation!")

        transform_list.append(transforms.ToTensor())

        if dataset in ["MNIST", "FashionMNIST", "EMNIST"]:
            transform_list = [
                transforms.Grayscale(3), ] + transform_list

        transform_ = transforms.Compose(transform_list)
        return transform_

    def get_data_loader(self):

        train_transform = self.get_data_transform(
            self.args.dataset, use_transform=None)