Add support for accimage.Image

It can be enabled by setting torchvision.image_backend = 'accimage'

Author: Maratyszcza

test/assets/grace_hopper_512.jpg

@@ -4,6 +4,7 @@
 from tqdm import tqdm
 import torch
 import torch.utils.data
+import torchvision
 import torchvision.transforms as transforms
 import torchvision.datasets as datasets
 
@@ -15,11 +16,17 @@
                     help='number of data loading threads (default: 2)')
 parser.add_argument('--batchSize', '-b', default=256, type=int, metavar='N',
                     help='mini-batch size (1 = pure stochastic) Default: 256')
+parser.add_argument('--accimage', action='store_true',
+                    help='use accimage')
 
 
 if __name__ == "__main__":
     args = parser.parse_args()
 
+    if args.accimage:
+        torchvision.image_backend = 'accimage'
+    print('Using {}'.format(torchvision.image_backend))
+
     # Data loading code
     transform = transforms.Compose([
         transforms.RandomSizedCrop(224),
@@ -38,11 +45,14 @@
     train_iter = iter(train_loader)
 
     start_time = timer()
-    batch_count = 100 * args.nThreads
-    for i in tqdm(xrange(batch_count)):
+    batch_count = 20 * args.nThreads
+    for _ in tqdm(range(batch_count)):
         batch = next(train_iter)
     end_time = timer()
-    print("Performance: {dataset:.0f} minutes/dataset, {batch:.2f} secs/batch, {image:.2f} ms/image".format(
-        dataset=(end_time - start_time) * (float(len(train_loader)) / batch_count / 60.0),
-        batch=(end_time - start_time) / float(batch_count),
-        image=(end_time - start_time) / (batch_count * args.batchSize) * 1.0e+3))
+    print("Performance: {dataset:.0f} minutes/dataset, {batch:.1f} ms/batch,"
+          " {image:.2f} ms/image {rate:.0f} images/sec"
+          .format(
+           dataset=(end_time - start_time) * (float(len(train_loader)) / batch_count / 60.0),
+           batch=(end_time - start_time) / float(batch_count) * 1.0e+3,
+           image=(end_time - start_time) / (batch_count * args.batchSize) * 1.0e+3,
+           rate=(batch_count * args.batchSize) / (end_time - start_time)))

test/preprocess-bench.py

@@ -3,6 +3,14 @@
 import unittest
 import random
 import numpy as np
+from PIL import Image
+try:
+    import accimage
+except ImportError:
+    accimage = None
+
+
+GRACE_HOPPER = 'assets/grace_hopper_512.jpg'
 
 
 class Tester(unittest.TestCase):
@@ -153,6 +161,51 @@ def test_to_tensor(self):
         expected_output = ndarray.transpose((2, 0, 1)) / 255.0
         assert np.allclose(output.numpy(), expected_output)
 
+    @unittest.skipIf(accimage is None, 'accimage not available')
+    def test_accimage_to_tensor(self):
+        trans = transforms.ToTensor()
+
+        expected_output = trans(Image.open(GRACE_HOPPER).convert('RGB'))
+        output = trans(accimage.Image(GRACE_HOPPER))
+
+        import visdom
+        vis = visdom.Visdom()
+        vis.image(expected_output.numpy())
+        vis.image(output.numpy())
+        vis.image((expected_output - output).numpy())
+
+        self.assertEqual(expected_output.size(), output.size())
+        assert np.allclose(output.numpy(), expected_output.numpy())
+
+    @unittest.skipIf(accimage is None, 'accimage not available')
+    def test_accimage_resize(self):
+        trans = transforms.Compose([
+            transforms.Scale(256, interpolation=Image.LINEAR),
+            transforms.ToTensor(),
+        ])
+
+        expected_output = trans(Image.open(GRACE_HOPPER).convert('RGB'))
+        output = trans(accimage.Image(GRACE_HOPPER))
+
+        self.assertEqual(expected_output.size(), output.size())
+        self.assertLess(np.abs((expected_output - output).mean()), 1e-3)
+        self.assertLess((expected_output - output).var(), 1e-5)
+        # note the high absolute tolerance
+        assert np.allclose(output.numpy(), expected_output.numpy(), atol=5e-2)
+
+    @unittest.skipIf(accimage is None, 'accimage not available')
+    def test_accimage_crop(self):
+        trans = transforms.Compose([
+            transforms.CenterCrop(256),
+            transforms.ToTensor(),
+        ])
+
+        expected_output = trans(Image.open(GRACE_HOPPER).convert('RGB'))
+        output = trans(accimage.Image(GRACE_HOPPER))
+
+        self.assertEqual(expected_output.size(), output.size())
+        assert np.allclose(output.numpy(), expected_output.numpy())
+
     def test_tensor_to_pil_image(self):
         trans = transforms.ToPILImage()
         to_tensor = transforms.ToTensor()

test/test_transforms.py

@@ -2,3 +2,17 @@
 from torchvision import datasets
 from torchvision import transforms
 from torchvision import utils
+
+
+image_backend = 'PIL'
+"""
+Specifies the package used to load images.
+
+Options are 'PIL' and 'accimage'. The :mod:`accimage` package uses the
+Intel IPP library. It is generally faster than PIL, but does not support as
+many operations.
+"""
+
+
+def get_image_backend():
+    return image_backend

torchvision/__init__.py

@@ -39,6 +39,14 @@ def make_dataset(dir, class_to_idx):
 
 
 def default_loader(path):
+    from torchvision import get_image_backend
+    if get_image_backend() == 'accimage':
+        import accimage
+        try:
+            return accimage.Image(path)
+        except IOError:
+            # Potentially a decoding problem, fall back to PIL.Image
+            pass
     return Image.open(path).convert('RGB')
 
 

torchvision/datasets/folder.py

@@ -3,6 +3,10 @@
 import math
 import random
 from PIL import Image, ImageOps
+try:
+    import accimage
+except ImportError:
+    accimage = None
 import numpy as np
 import numbers
 import types
@@ -42,6 +46,12 @@ def __call__(self, pic):
             img = torch.from_numpy(pic.transpose((2, 0, 1)))
             # backard compability
             return img.float().div(255)
+
+        if accimage is not None and isinstance(pic, accimage.Image):
+            nppic = np.zeros([pic.channels, pic.height, pic.width], dtype=np.float32)
+            pic.copyto(nppic)
+            return torch.from_numpy(nppic)
+
         # handle PIL Image
         if pic.mode == 'I':
             img = torch.from_numpy(np.array(pic, np.int32, copy=False))