Merge pull request #4 from pytorch/numpy

fix ToTensor to handle numpy

Author: soumith

test/cifar.py

@@ -1,12 +1,35 @@
 import torch
 import torchvision.datasets as dset
+import torchvision.transforms as transforms
 
-print('\n\nCifar 10')
-a = dset.CIFAR10(root="abc/def/ghi", download=True)
+# print('\n\nCifar 10')
+# a = dset.CIFAR10(root="abc/def/ghi", download=True)
 
-print(a[3])
+# print(a[3])
 
-print('\n\nCifar 100')
-a = dset.CIFAR100(root="abc/def/ghi", download=True)
+# print('\n\nCifar 100')
+# a = dset.CIFAR100(root="abc/def/ghi", download=True)
 
-print(a[3])
+# print(a[3])
+
+
+dataset = dset.CIFAR10(root='cifar', download=True, transform=transforms.ToTensor())
+dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, 
+                                         shuffle=True, num_workers=2)
+
+
+# miter = dataloader.__iter__()
+# def getBatch():
+#     global miter
+#     try:
+#         return miter.next()
+#     except StopIteration:
+#         miter = dataloader.__iter__()
+#         return miter.next()
+    
+# i=0
+# while True:
+#     print(i)
+#     img, target = getBatch()
+#     i+=1
+    

torchvision/transforms.py

@@ -2,6 +2,7 @@
 import math
 import random
 from PIL import Image
+import numpy as np
 
 
 class Compose(object):
@@ -16,11 +17,16 @@ def __call__(self, img):
 
 class ToTensor(object):
     def __call__(self, pic):
-        img = torch.ByteTensor(torch.ByteStorage.from_buffer(pic.tobytes()))
-        img = img.view(pic.size[0], pic.size[1], 3)
-        # put it in CHW format
-        # yikes, this transpose takes 80% of the loading time/CPU
-        img = img.transpose(0, 2).transpose(1, 2).contiguous()
+        if isinstance(pic, np.ndarray):
+            # handle numpy array
+            img = torch.from_numpy(pic)
+        else:
+            # handle PIL Image
+            img = torch.ByteTensor(torch.ByteStorage.from_buffer(pic.tobytes()))
+            img = img.view(pic.size[0], pic.size[1], 3)
+            # put it in CHW format
+            # yikes, this transpose takes 80% of the loading time/CPU
+            img = img.transpose(0, 2).transpose(1, 2).contiguous()
         return img.float()
 
 class Normalize(object):