A guide on Colab TPU training using PyTorch XLA (Part 5.1)

dataset_dwnld_dir.py

PATH = "./my_dataset/" # path to dataset on Colab instance
TRAIN_PATH = PATH + "train/"
VAL_PATH = PATH + "val/"

# your custom augmentations
T = transforms.Compose([
      transforms.ToTensor(),
      ...
    ])

# same procedure but you use datasets.ImageFolder(...) instead
if not xm.is_master_ordinal():
  xm.rendezvous('download_only_once')

data_train = datasets.ImageFolder(TRAIN_PATH, transform=T)
data_test = datasets.ImageFolder(VAL_PATH, transform=T)

if xm.is_master_ordinal():
  xm.rendezvous('download_only_once')