PyTorch/CUDA interoperability

Overview

This sample shows how interoperability between VPI and PyTorch works in Python. It shows a simple set of image processing operations applied on a sample input image in Pytorch and VPI without any memory copies during the interoperation.

This sample uses PyTorch as an example, however, this interoperability is possible with any library that supports the cuda array interface.

Libraries such as:

• Numba
• CuPy
• PyArrow
• mpi4py
• ArrayViews
• JAX
• PyCUDA
• NVIDIA DALI

The sample starts with an image given by you, converted to grayscale, then performs the following sequence of operations to it:

1. Lower the image intensity using PyTorch
2. Does PyTorch -> VPI interoperability without memory copies

   vpi_image = vpi.asimage(torch_image)

3. Applies a box filter using VPI
4. Does VPI -> PyTorch interoperability without memory copies

   with vpi_image.rlock_cuda() as cuda_buffer:
       torch_image = torch.as_tensor(cuda_buffer)

   or, with an additional deep memory copy

   torch_image = torch.as_tensor(vpi_image.cuda())

5. Increase the image intensity using PyTorch

The result is then saved to disk.

Instructions

The usage is:

python3 main.py <input image>

where

• input image: input image file name; it accepts png, jpeg and possibly others.

Here's one example:

• Python

  python3 main.py ../assets/kodim08.png

The sample will produce vpi_pytorch.png as output file.

Results

Input image	Output image

Source Code

For convenience, here's the code that is also installed in the samples directory.

import vpi
import torch
import numpy as np
 
from PIL import Image, ImageOps
from argparse import ArgumentParser
 
# Parse command line arguments
parser = ArgumentParser()
parser.add_argument('input', help='Image to be used as input')
 
args = parser.parse_args()
 
# Make sure CUDA is available for this example
assert torch.cuda.is_available()
cuda_device = torch.device('cuda')
 
# Read the input image and convert it to grayscale
try:
    pil_image = ImageOps.grayscale(Image.open(args.input))
except IOError:
    sys.exit("Input file not found")
except:
    sys.exit("Error with input file")
 
# Pillow -> NumPy --------------------------------------
np_image = np.asarray(pil_image)
 
# NumPy -> PyTorch/CUDA --------------------------------
torch_image = torch.asarray(np_image).cuda()
 
# Perform an operation using PyTorch
torch_image = torch_image/255 * 0.5
 
# PyTorch/CUDA -> VPI, no copies involved --------------
vpi_image = vpi.asimage(torch_image)
 
# Peform operations using VPI's CUDA backend
with vpi.Backend.CUDA:
    # Blur the input image with box filter
    vpi_image = vpi_image.box_filter(3)
 
# VPI -> PyTorch/CUDA, no copies involved --------------
with vpi_image.rlock_cuda() as cuda_buffer:
    # Perform another operation using PyTorch
    torch_tensor = torch.as_tensor(cuda_buffer, device=cuda_device)
    torch_image = torch_tensor*255 + 128
 
# PyTorch -> NumPy -------------------------------------
np_image = np.asarray(torch_image.cpu())
 
# NumPy -> Pillow --------------------------------------
pil_image = Image.fromarray(np_image)
 
# Save the output to disk
pil_image.convert('L').save('vpi_pytorch.png')