FedML-AI/FedML

How to change data transform on built-in dataset?

HeywardLiu opened this issue · 0 comments

I've read how you implement cifar10 in

  • fedml/data/data_loader.py : Entry point of loading dataset
  • fedml/data/cifar10/efficient_loader.py : loading cifar10 efficiently

I'm wondering if it is possible to change _data_transforms to a customized one? Can I configure it in fedml_config.yaml?

Since I want to perform federated transfer learning (FTL) from the ImageNet-pretrained model,
which have been trained at different resolution (224x224).

Thanks!

def _data_transforms_cifar10():
    CIFAR_MEAN = [0.49139968, 0.48215827, 0.44653124]
    CIFAR_STD = [0.24703233, 0.24348505, 0.26158768]

    train_transform = transforms.Compose(
        [
            transforms.ToPILImage(),
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
        ]
    )

    train_transform.transforms.append(Cutout(16))

    valid_transform = transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
        ]
    )

    return train_transform, valid_transform


def load_cifar10_data(datadir):
    train_transform, test_transform = _data_transforms_cifar10()

    cifar10_train_ds = CIFAR10_truncated(
        datadir, train=True, download=True, transform=train_transform
    )
    cifar10_test_ds = CIFAR10_truncated(
        datadir, train=False, download=True, transform=test_transform
    )

    X_train, y_train = cifar10_train_ds.data, cifar10_train_ds.target
    X_test, y_test = cifar10_test_ds.data, cifar10_test_ds.target

    return (X_train, y_train, X_test, y_test)