CNN / src /data_loader.py
Krippa's picture
Upload folder using huggingface_hub
ecf179e verified
Raw
History Blame Contribute Delete
1.18 kB
import torch
import torchvision
import torchvision.transforms as transforms
def load_data(batch_size=64):
"""
Loads and preprocesses the CIFAR-10 dataset using PyTorch.
Returns:
tuple: (trainloader, testloader)
"""
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
shuffle=True, num_workers=0)
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,
shuffle=False, num_workers=0)
classes = ('plane', 'car', 'bird', 'cat',
'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
return trainloader, testloader, classes