def imagenet_train(dataset_root, nb_classes, patch_size, patch_overlap):
train_transform = amdim_transforms.TransformsImageNet128Patches(
patch_size=patch_size, overlap=patch_overlap)
dataset = UnlabeledImagenet(dataset_root,
nb_classes=nb_classes,
split='train',
transform=train_transform)
return dataset
def imagenet(dataset_root, nb_classes, patch_size, patch_overlap, split: str = 'train'):
assert split in ('train', 'val')
train_transform = amdim_transforms.TransformsImageNet128Patches(
patch_size=patch_size,
overlap=patch_overlap
)
dataset = UnlabeledImagenet(
dataset_root,
nb_classes=nb_classes,
split=split,
transform=train_transform,
)
return dataset