{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Processing GPU Data with Python Operators" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "This example shows you how to use the `PythonFunction` operator on a GPU. For an introduction and general information about Python operators family see the [Python Operators](./python_operator.ipynb) section.\n", "\n", "Although Python operators are not designed to be fast, it might be useful to run them on a GPU, for example, when we want to introduce a custom operation to an existing GPU pipeline. For this purpose, all operators in the `PythonFunction` family have their GPU variants.\n", "\n", "For the `TorchPythonFunction` and `DLTensorPythonFunction` operators data formats on which they operate stay the same as for the CPU, PyTorch tensors in the former, and [DLPack](https://github.com/dmlc/dlpack) tensors in the latter. For the GPU `PythonFunction`, inputs and outputs of the implementation function are [CuPy](https://cupy.chainer.org/) arrays." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## CuPy Operations\n", "Since the *CuPy* arrays API is similar to the API in NumPy, we can implement the same operation that we defined in the CPU example nearly without any code changes. " ] }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ "from nvidia.dali.pipeline import Pipeline\n", "import nvidia.dali.fn as fn\n", "import nvidia.dali.types as types\n", "import numpy\n", "import cupy\n", "\n", "\n", "def edit_images(image1, image2):\n", " assert image1.shape == image2.shape\n", " h, w, c = image1.shape\n", " y, x = cupy.ogrid[0:h, 0:w]\n", " mask = (x - w / 2) ** 2 + (y - h / 2) ** 2 > h * w / 9\n", " result1 = cupy.copy(image1)\n", " result1[mask] = image2[mask]\n", " result2 = cupy.copy(image2)\n", " result2[mask] = image1[mask]\n", " return result1, result2" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Another way to define a GPU function with *CuPy* is to write a CUDA kernel. Here, we present a simple kernel that interleaves channels of the two images. Refer to the *CuPy* [documentation](https://docs-cupy.chainer.org/en/stable/tutorial/kernel.html) for more information." ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "mix_channels_kernel = cupy.ElementwiseKernel(\n", " \"uint8 x, uint8 y\", \"uint8 z\", \"z = (i % 3) ? x : y\", \"mix_channels\"\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "