CIFAR10 dataloader

dataset/loaders.py

@@ -3,6 +3,8 @@
 import numpy as np
 import torch
 import random
+import torchvision
+import torchvision.transforms as transforms
 from torch.autograd import Variable
 
 use_cuda = True # maybe change this to infer the device
@@ -33,4 +35,25 @@ def get_8gaussians(batch_size):
         out = Variable(torch.Tensor(dataset))
         if use_cuda:
             out = out.cuda()
-        yield out
+        yield out
+
+def CIFAR(batch_size = 64):
+    """The function returns the train and test loaders as well as the class labels in human readable form"""
+    transform = transforms.Compose(
+    [transforms.ToTensor(),
+     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
+
+    train_set = torchvision.datasets.CIFAR10(root='./data', train=True,
+                                        download=True, transform=transform)
+    train_loader = torch.utils.data.DataLoader(train_set, batch_size=batch_size,
+                                          shuffle=True, num_workers=2)
+
+    test_set = torchvision.datasets.CIFAR10(root='./data', train=False,
+                                       download=True, transform=transform)
+    test_loader = torch.utils.data.DataLoader(test_set, batch_size=batch_size,
+                                         shuffle=False, num_workers=2)
+
+    classes = ('plane', 'car', 'bird', 'cat',
+           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
+
+    return train_loader, test_loader, classes