#number of sub-processes
num_workers = 0
#how many samples per batch to load
batch_size = 40

#transform for the train dataset with aumuntation
transform_trn = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.RandomHorizontalFlip(), #randomly flip and rotate
    transforms.RandomRotation(10),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])
    ])

#transform for validation and testing datasets
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])
transform = transforms.Compose([
           transforms.Resize(256), 
            transforms.ToTensor(),
            normalize,
        ])

#choose the training and test datasets
train_data = datasets.ImageFolder('/data/dog_images/train', transform=transform_trn)
val_data = datasets.ImageFolder('/data/dog_images/valid', transform=transform)
test_data = datasets.ImageFolder('/data/dog_images/test', transform=transform)

train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size,
    shuffle=True, num_workers=num_workers)
valid_loader = torch.utils.data.DataLoader(val_data, batch_size=batch_size, 
    shuffle=False, num_workers=num_workers)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size,shuffle=False,
    num_workers=num_workers)