INNER CODE UNIT · Python

get_dataloader

digantamisra98/Mish · PyTorch Benchmarks/train_imagenet.py:209

def get_dataloader(root_dir, is_train, batch_size, workers):
    dir_name = "train" if is_train else "val"
    data_dir = os.path.join(root_dir, dir_name)
    normalize = transforms.Normalize(
        mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
    )
    transformations = (
        [
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
        ]
        if is_train
        else [
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),

View source record →