Training with CELEB-A

[AGKhalil]

Hello and thank you for the great repo. In here it is mentioned:

Note The data in the train_data.npz and eval_data.npz files must be loadable as follows:

train_data = np.load(os.path.join(PATH, f'data/{args.dataset}', 'train_data.npz'))['data']
eval_data = np.load(os.path.join(PATH, f'data/{args.dataset}', 'eval_data.npz'))['data']

This works with MNIST and CIFAR-10, but CELEB-A is simply too large to load from a single npz file. Am I approaching this correctly?

[clementchadebec]

Hello @AGKhalil,

Thank you very much for the kind words and your interest in the repo.

You can indeed either load everything in RAM if your compute allows it or better, you can create your own Dataset (that will for instance load the .jpg images on the fly) and pass it to the Trainer instance. You can have a look to this example on ImageNet dataset or the summary below extracted from this example. By the way, it also shows how to perform distributed training which may be useful for those large databases.

from PIL import Image
from torch.utils.data import Dataset

from pythae.data.datasets import DatasetOutput


# Create your dataset
class ImageNet(Dataset):
    def __init__(self, data_dir=None, transforms=None):
        self.imgs_path = [os.path.join(data_dir, n) for n in os.listdir(data_dir)]
        self.transforms = transforms

    def __len__(self):
        return len(self.imgs_path)

    def __getitem__(self, idx):
        img = Image.open(self.imgs_path[idx]).convert("RGB")
        if self.transforms is not None:
            img = self.transforms(img)
        return DatasetOutput(data=img)

# define your pre-processing
img_transforms = transforms.Compose(
        [transforms.Resize((128, 128)), transforms.ToTensor()]
)

# instantiate the datasets
train_dataset = ImageNet(
        data_dir="/gpfsscratch/rech/wlr/uhw48em/data/imagenet/train",
        transforms=img_transforms,
)
eval_dataset = ImageNet(
      data_dir="/gpfsscratch/rech/wlr/uhw48em/data/imagenet/val",
      transforms=img_transforms,
)

# pass them to the Trainer
trainer = BaseTrainer(
        model=model,
        train_dataset=train_dataset, ### here
        eval_dataset=eval_dataset, ### here
        training_config=training_config,
        callbacks=callbacks,
)

I hope this helps :)

Best,

Clément