PyTorch/CUDA 互操作性

概述

此示例展示了 VPI 和 PyTorch 之间在 Python 中的互操作性如何工作。它展示了一组简单的图像处理操作，应用于 PyTorch 和 VPI 中的示例输入图像，而无需在互操作期间进行任何内存复制。

此示例使用 PyTorch 作为示例，但是，这种互操作性可以与任何支持 cuda 数组接口的库一起使用。

诸如以下的库

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

该示例从您提供的图像开始，将其转换为灰度图像，然后对其执行以下操作序列

1. 使用 PyTorch 降低图像强度
2. 执行 PyTorch -> VPI 互操作性，无需内存复制

   vpi_image = vpi.asimage(torch_image)

3. 使用 VPI 应用盒状滤波器
4. 执行 VPI -> PyTorch 互操作性，无需内存复制

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

   或者，使用额外的深度内存复制

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

5. 使用 PyTorch 增加图像强度

然后将结果保存到磁盘。

说明

用法是

python3 main.py <输入图像>

其中

• 输入图像：输入图像文件名；它接受 png、jpeg 以及可能的其他格式。

这是一个示例

• Python

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

该示例将生成 vpi_pytorch.png 作为输出文件。

结果

输入图像	输出图像

源代码

为了方便起见，以下代码也安装在示例目录中。

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')