20 lines
850 B
Python
20 lines
850 B
Python
import torchvision
|
|
import torchvision.transforms as transforms
|
|
import torch.utils.data
|
|
|
|
def get_data_loaders(data_dir, batch_size):
|
|
transform = transforms.Compose([
|
|
transforms.RandomHorizontalFlip(),
|
|
transforms.RandomCrop(32, padding=4),
|
|
transforms.ToTensor(),
|
|
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
|
|
])
|
|
|
|
train_dataset = torchvision.datasets.CIFAR10(root=data_dir, train=True, download=True, transform=transform)
|
|
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
|
|
|
|
test_dataset = torchvision.datasets.CIFAR10(root=data_dir, train=False, download=True, transform=transform)
|
|
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4)
|
|
|
|
return train_loader, test_loader
|