diff --git a/test/assets/grace_hopper_517x606.jpg b/test/assets/grace_hopper_517x606.jpg new file mode 100644 index 00000000000..d2a427810f6 Binary files /dev/null and b/test/assets/grace_hopper_517x606.jpg differ diff --git a/test/preprocess-bench.py b/test/preprocess-bench.py index 7746c9c3111..55bed4ef53a 100644 --- a/test/preprocess-bench.py +++ b/test/preprocess-bench.py @@ -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.set_image_backend('accimage') + print('Using {}'.format(torchvision.get_image_backend())) + # Data loading code transform = transforms.Compose([ transforms.RandomSizedCrop(224), @@ -38,11 +45,13 @@ 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))) diff --git a/test/test_transforms.py b/test/test_transforms.py index 64327ac0584..f14c703a603 100644 --- a/test/test_transforms.py +++ b/test/test_transforms.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_517x606.jpg' class Tester(unittest.TestCase): @@ -153,6 +161,45 @@ 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)) + + 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() diff --git a/torchvision/__init__.py b/torchvision/__init__.py index 5f8137910b4..e31b8e998b0 100644 --- a/torchvision/__init__.py +++ b/torchvision/__init__.py @@ -2,3 +2,31 @@ from torchvision import datasets from torchvision import transforms from torchvision import utils + + +_image_backend = 'PIL' + + +def set_image_backend(backend): + """ + 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. + + Args: + backend (string): name of the image backend + """ + global _image_backend + if backend not in ['PIL', 'accimage']: + raise ValueError("Invalid backend '{}'. Options are 'PIL' and 'accimage'" + .format(backend)) + _image_backend = backend + + +def get_image_backend(): + """ + Gets the name of the package used to load images + """ + return _image_backend diff --git a/torchvision/datasets/folder.py b/torchvision/datasets/folder.py index 633ebab0153..3f15c647103 100644 --- a/torchvision/datasets/folder.py +++ b/torchvision/datasets/folder.py @@ -38,10 +38,27 @@ def make_dataset(dir, class_to_idx): return images -def default_loader(path): +def pil_loader(path): return Image.open(path).convert('RGB') +def accimage_loader(path): + import accimage + try: + return accimage.Image(path) + except IOError: + # Potentially a decoding problem, fall back to PIL.Image + return pil_loader(path) + + +def default_loader(path): + from torchvision import get_image_backend + if get_image_backend() == 'accimage': + return accimage_loader(path) + else: + return pil_loader(path) + + class ImageFolder(data.Dataset): def __init__(self, root, transform=None, target_transform=None, diff --git a/torchvision/transforms.py b/torchvision/transforms.py index 6ae1e79beaf..52a15a0bdd9 100644 --- a/torchvision/transforms.py +++ b/torchvision/transforms.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))