advanced/python_custom_ops

在 Google Colab 中运行

Colab

下载 Notebook

Notebook

在 GitHub 上查看

GitHub

自定义 Python 算子

创建日期：2024 年 6 月 18 日 | 最后更新：2025 年 3 月 19 日 | 最后验证：2024 年 11 月 5 日

您将学到什么

• 如何将使用 Python 编写的自定义算子集成到 PyTorch 中

• 如何使用 torch.library.opcheck 测试自定义算子

先决条件

• PyTorch 2.4 或更高版本

PyTorch 提供了大量的算子库来处理张量（例如 torch.add, torch.sum 等）。然而，您可能希望在 PyTorch 中使用新的自定义算子，例如由第三方库提供的算子。本教程将展示如何包装 Python 函数，使其行为如同 PyTorch 原生算子。您可能希望在 PyTorch 中创建自定义算子的原因包括：

• 将任意 Python 函数视为对 torch.compile 不透明的可调用对象（即，防止 torch.compile 对该函数进行追踪）。

• 为任意 Python 函数添加训练支持。

使用 torch.library.custom_op() 创建 Python 自定义算子。使用 C++ TORCH_LIBRARY API 创建 C++ 自定义算子（这些算子可在无 Python 的环境中使用）。更多详细信息，请参阅 自定义算子着陆页。

请注意，如果您的操作可以表示为现有 PyTorch 算子的组合，那么通常不需要使用自定义算子 API —— 一切（例如 torch.compile，训练支持）都应该可以直接工作。

示例：将 PIL 的 crop 包装为自定义算子

假设我们正在使用 PIL 的 crop 操作。

import torch
from torchvision.transforms.functional import to_pil_image, pil_to_tensor
import PIL
import IPython
import matplotlib.pyplot as plt

def crop(pic, box):
    img = to_pil_image(pic.cpu())
    cropped_img = img.crop(box)
    return pil_to_tensor(cropped_img).to(pic.device) / 255.

def display(img):
    plt.imshow(img.numpy().transpose((1, 2, 0)))

img = torch.ones(3, 64, 64)
img *= torch.linspace(0, 1, steps=64) * torch.linspace(0, 1, steps=64).unsqueeze(-1)
display(img)

cropped_img = crop(img, (10, 10, 50, 50))
display(cropped_img)

crop 不能被 torch.compile 有效地直接处理：torch.compile 会在其无法处理的函数上触发 “图形中断 (graph break)”，而图形中断会对性能产生不利影响。以下代码通过抛出错误演示了这一点（如果发生图形中断，torch.compile 设置 fullgraph=True 时会抛出错误）。

@torch.compile(fullgraph=True)
def f(img):
    return crop(img, (10, 10, 50, 50))

# The following raises an error. Uncomment the line to see it.
# cropped_img = f(img)

为了将 crop 作为黑盒处理以供 torch.compile 使用，我们需要做两件事：

1. 将该函数包装为 PyTorch 自定义算子。

2. 为该算子添加一个 “FakeTensor 内核”（也称为 “meta 内核”）。给定一些 FakeTensors 输入（不带存储的虚拟张量），该函数应返回您选择的、具有正确张量元数据（形状/步长/dtype/设备）的虚拟张量。

from typing import Sequence

# Use torch.library.custom_op to define a new custom operator.
# If your operator mutates any input Tensors, their names must be specified
# in the ``mutates_args`` argument.
@torch.library.custom_op("mylib::crop", mutates_args=())
def crop(pic: torch.Tensor, box: Sequence[int]) -> torch.Tensor:
    img = to_pil_image(pic.cpu())
    cropped_img = img.crop(box)
    return (pil_to_tensor(cropped_img) / 255.).to(pic.device, pic.dtype)

# Use register_fake to add a ``FakeTensor`` kernel for the operator
@crop.register_fake
def _(pic, box):
    channels = pic.shape[0]
    x0, y0, x1, y1 = box
    result = pic.new_empty(y1 - y0, x1 - x0, channels).permute(2, 0, 1)
    # The result should have the same metadata (shape/strides/``dtype``/device)
    # as running the ``crop`` function above.
    return result

在此之后，crop 就可以在没有图形中断的情况下工作了。

@torch.compile(fullgraph=True)
def f(img):
    return crop(img, (10, 10, 50, 50))

cropped_img = f(img)
display(img)

display(cropped_img)

为 crop 添加训练支持

使用 torch.library.register_autograd 为算子添加训练支持。优先使用此方法，而不是直接使用 torch.autograd.Function；当 autograd.Function 与 PyTorch 算子注册 API 组合使用时，在与 torch.compile 配合时可能会（且曾经）导致隐蔽的不正确性。

如果您不需要训练支持，则无需使用 torch.library.register_autograd。如果您最终使用没有 autograd 注册的 custom_op 进行训练，我们将抛出一条错误消息。

crop 的梯度公式本质上是 PIL.paste（我们将推导过程留给读者作为练习）。首先让我们将 paste 包装成一个自定义算子。

@torch.library.custom_op("mylib::paste", mutates_args=())
def paste(im1: torch.Tensor, im2: torch.Tensor, coord: Sequence[int]) -> torch.Tensor:
    assert im1.device == im2.device
    assert im1.dtype == im2.dtype
    im1_pil = to_pil_image(im1.cpu())
    im2_pil = to_pil_image(im2.cpu())
    PIL.Image.Image.paste(im1_pil, im2_pil, coord)
    return (pil_to_tensor(im1_pil) / 255.).to(im1.device, im1.dtype)

@paste.register_fake
def _(im1, im2, coord):
    assert im1.device == im2.device
    assert im1.dtype == im2.dtype
    return torch.empty_like(im1)

现在，让我们使用 register_autograd 来指定 crop 的梯度公式。

def backward(ctx, grad_output):
    grad_input = grad_output.new_zeros(ctx.pic_shape)
    grad_input = paste(grad_input, grad_output, ctx.coords)
    return grad_input, None

def setup_context(ctx, inputs, output):
    pic, box = inputs
    ctx.coords = box[:2]
    ctx.pic_shape = pic.shape

crop.register_autograd(backward, setup_context=setup_context)

请注意，反向传播（backward）必须是 PyTorch 可理解算子的组合，这就是为什么我们将 paste 包装成自定义算子，而不是直接使用 PIL 的 paste 的原因。

img = img.requires_grad_()
result = crop(img, (10, 10, 50, 50))
result.sum().backward()
display(img.grad)

这是正确的梯度，裁剪区域显示为 1（白色），未使用的区域显示为 0（黑色）。

测试 Python 自定义算子

使用 torch.library.opcheck 测试自定义算子是否注册正确。这不会测试梯度在数学上是否正确；请为此编写单独的测试（手动测试或使用 torch.autograd.gradcheck）。

要使用 opcheck，请向其传递一组示例输入进行测试。如果您的算子支持训练，那么示例应包含需要梯度的张量。如果您的算子支持多种设备，那么示例应包含来自每个设备的张量。

examples = [
    [torch.randn(3, 64, 64), [0, 0, 10, 10]],
    [torch.randn(3, 91, 91, requires_grad=True), [10, 0, 20, 10]],
    [torch.randn(3, 60, 60, dtype=torch.double), [3, 4, 32, 20]],
    [torch.randn(3, 512, 512, requires_grad=True, dtype=torch.double), [3, 4, 32, 45]],
]

for example in examples:
    torch.library.opcheck(crop, example)

可变（Mutable）Python 自定义算子

您也可以将修改其输入的 Python 函数包装为自定义算子。修改输入的函数很常见，因为许多底层内核就是这样编写的；例如，计算 sin 的内核可能会接收输入和一个输出张量，并将 input.sin() 写入输出张量。

我们将使用 numpy.sin 来演示可变 Python 自定义算子的示例。

import numpy as np

@torch.library.custom_op("mylib::numpy_sin", mutates_args={"output"}, device_types="cpu")
def numpy_sin(input: torch.Tensor, output: torch.Tensor) -> None:
    assert input.device == output.device
    assert input.device.type == "cpu"
    input_np = input.numpy()
    output_np = output.numpy()
    np.sin(input_np, out=output_np)

因为该算子不返回任何内容，所以无需注册 FakeTensor 内核（meta 内核）即可使其与 torch.compile 一起工作。

@torch.compile(fullgraph=True)
def f(x):
    out = torch.empty(3)
    numpy_sin(x, out)
    return out

x = torch.randn(3)
y = f(x)
assert torch.allclose(y, x.sin())

这是 opcheck 的运行结果，告知我们确实正确注册了算子。例如，如果我们忘记将输出添加到 mutates_args 中，opcheck 将报错。

example_inputs = [
    [torch.randn(3), torch.empty(3)],
    [torch.randn(0, 3), torch.empty(0, 3)],
    [torch.randn(1, 2, 3, 4, dtype=torch.double), torch.empty(1, 2, 3, 4, dtype=torch.double)],
]

for example in example_inputs:
    torch.library.opcheck(numpy_sin, example)

结论

在本教程中，我们学习了如何使用 torch.library.custom_op 在 Python 中创建与 PyTorch 子系统（如 torch.compile 和自动求导）协同工作的自定义算子。

本教程提供了自定义算子的基本介绍。有关更多详细信息，请参阅：

• torch.library 文档

• 自定义算子手册