在 torch.compile/torch.export 中支持自定义 C++ 类#

本教程是自定义 C++ 类教程的后续，介绍了在 torch.compile/torch.export 中支持自定义 C++ 类所需的额外步骤。

警告

此功能处于原型状态，可能会发生向后不兼容的更改。本教程提供了截至 PyTorch 2.8 的快照。如果您遇到任何问题，请在 Github 上与我们联系！

具体来说，有以下几个步骤：

在 C++ 自定义类实现中实现一个 __obj_flatten__ 方法，以允许我们检查其状态并保护更改。该方法应返回一个由属性名和值组成的元组的元组 (tuple[tuple[str, value] * n])。
使用 @torch._library.register_fake_class 注册一个 Python 伪类。
1. 为该类的每个 C++ 方法实现“伪方法”，这些伪方法应具有与 C++ 实现相同的模式。
2. 此外，在 Python 伪类中实现一个 __obj_unflatten__ 类方法，告诉我们如何从 __obj_flatten__ 返回的扁平化状态创建一个伪类。

以下是差异的分解。按照使用自定义 C++ 类扩展 TorchScript中的指南，我们可以创建一个线程安全的张量队列并构建它。

// Thread-safe Tensor Queue

#include <torch/custom_class.h>
#include <torch/script.h>

#include <iostream>
#include <string>
#include <vector>

using namespace torch::jit;

// Thread-safe Tensor Queue
struct TensorQueue : torch::CustomClassHolder {
explicit TensorQueue(at::Tensor t) : init_tensor_(t) {}

explicit TensorQueue(c10::Dict<std::string, at::Tensor> dict) {
    init_tensor_ = dict.at(std::string("init_tensor"));
    const std::string key = "queue";
    at::Tensor size_tensor;
    size_tensor = dict.at(std::string(key + "/size")).cpu();
    const auto* size_tensor_acc = size_tensor.const_data_ptr<int64_t>();
    int64_t queue_size = size_tensor_acc[0];

    for (const auto index : c10::irange(queue_size)) {
        at::Tensor val;
        queue_[index] = dict.at(key + "/" + std::to_string(index));
        queue_.push_back(val);
    }
}

// Push the element to the rear of queue.
// Lock is added for thread safe.
void push(at::Tensor x) {
    std::lock_guard<std::mutex> guard(mutex_);
    queue_.push_back(x);
}
// Pop the front element of queue and return it.
// If empty, return init_tensor_.
// Lock is added for thread safe.
at::Tensor pop() {
    std::lock_guard<std::mutex> guard(mutex_);
    if (!queue_.empty()) {
        auto val = queue_.front();
        queue_.pop_front();
        return val;
    } else {
        return init_tensor_;
    }
}

std::vector<at::Tensor> get_raw_queue() {
    std::vector<at::Tensor> raw_queue(queue_.begin(), queue_.end());
    return raw_queue;
}

private:
    std::deque<at::Tensor> queue_;
    std::mutex mutex_;
    at::Tensor init_tensor_;
};

// The torch binding code
TORCH_LIBRARY(MyCustomClass, m) {
    m.class_<TensorQueue>("TensorQueue")
        .def(torch::init<at::Tensor>())
        .def("push", &TensorQueue::push)
        .def("pop", &TensorQueue::pop)
        .def("get_raw_queue", &TensorQueue::get_raw_queue);
}

步骤 1：向 C++ 自定义类实现添加一个 __obj_flatten__ 方法

// Thread-safe Tensor Queue
struct TensorQueue : torch::CustomClassHolder {
...
std::tuple<std::tuple<std::string, std::vector<at::Tensor>>, std::tuple<std::string, at::Tensor>> __obj_flatten__() {
    return std::tuple(std::tuple("queue", this->get_raw_queue()), std::tuple("init_tensor_", this->init_tensor_.clone()));
}
...
};

TORCH_LIBRARY(MyCustomClass, m) {
    m.class_<TensorQueue>("TensorQueue")
        .def(torch::init<at::Tensor>())
        ...
        .def("__obj_flatten__", &TensorQueue::__obj_flatten__);
}

步骤 2a：在 Python 中注册一个实现每个方法的伪类。

# namespace::class_name
@torch._library.register_fake_class("MyCustomClass::TensorQueue")
class FakeTensorQueue:
    def __init__(
        self,
        queue: List[torch.Tensor],
        init_tensor_: torch.Tensor
    ) -> None:
        self.queue = queue
        self.init_tensor_ = init_tensor_

    def push(self, tensor: torch.Tensor) -> None:
        self.queue.append(tensor)

    def pop(self) -> torch.Tensor:
        if len(self.queue) > 0:
            return self.queue.pop(0)
        return self.init_tensor_

步骤 2b：在 Python 中实现一个 __obj_unflatten__ 类方法。

# namespace::class_name
@torch._library.register_fake_class("MyCustomClass::TensorQueue")
class FakeTensorQueue:
    ...
    @classmethod
    def __obj_unflatten__(cls, flattened_tq):
        return cls(**dict(flattened_tq))

就是这样！现在我们可以创建一个使用此对象的模块，并使用 torch.compile 或 torch.export 运行它。

import torch

torch.classes.load_library("build/libcustom_class.so")
tq = torch.classes.MyCustomClass.TensorQueue(torch.empty(0).fill_(-1))

class Mod(torch.nn.Module):
    def forward(self, tq, x):
        tq.push(x.sin())
        tq.push(x.cos())
        poped_t = tq.pop()
        assert torch.allclose(poped_t, x.sin())
        return tq, poped_t

tq, poped_t = torch.compile(Mod(), backend="eager", fullgraph=True)(tq, torch.randn(2, 3))
assert tq.size() == 1

exported_program = torch.export.export(Mod(), (tq, torch.randn(2, 3),), strict=False)
exported_program.module()(tq, torch.randn(2, 3))

我们还可以实现接受自定义类作为输入的自定义操作。例如，我们可以注册一个自定义操作 for_each_add_(tq, tensor)

struct TensorQueue : torch::CustomClassHolder {
    ...
    void for_each_add_(at::Tensor inc) {
        for (auto& t : queue_) {
            t.add_(inc);
        }
    }
    ...
}


TORCH_LIBRARY_FRAGMENT(MyCustomClass, m) {
    m.class_<TensorQueue>("TensorQueue")
        ...
        .def("for_each_add_", &TensorQueue::for_each_add_);

    m.def(
        "for_each_add_(__torch__.torch.classes.MyCustomClass.TensorQueue foo, Tensor inc) -> ()");
}

void for_each_add_(c10::intrusive_ptr<TensorQueue> tq, at::Tensor inc) {
    tq->for_each_add_(inc);
}

TORCH_LIBRARY_IMPL(MyCustomClass, CPU, m) {
    m.impl("for_each_add_", for_each_add_);
}

由于伪类是在 Python 中实现的，我们要求自定义操作的伪实现也必须在 Python 中注册。

@torch.library.register_fake("MyCustomClass::for_each_add_")
def fake_for_each_add_(tq, inc):
    tq.for_each_add_(inc)

重新编译后，我们可以使用以下命令导出自定义操作：

class ForEachAdd(torch.nn.Module):
    def forward(self, tq: torch.ScriptObject, a: torch.Tensor) -> torch.ScriptObject:
        torch.ops.MyCustomClass.for_each_add_(tq, a)
        return tq

mod = ForEachAdd()
tq = empty_tensor_queue()
qlen = 10
for i in range(qlen):
    tq.push(torch.zeros(1))

ep = torch.export.export(mod, (tq, torch.ones(1)), strict=False)

为什么我们需要创建一个伪类？#

使用真实自定义对象进行追踪有几个主要缺点：

对真实对象的操作可能很耗时，例如，自定义对象可能正在从网络读取数据或从磁盘加载数据。
我们不希望在追踪过程中改变真实的自定义对象或对环境产生副作用。
它不支持动态形状。

然而，对于用户来说，编写一个伪类可能很困难，例如，如果原始类使用了一些确定方法输出形状的第三方库，或者它很复杂且由他人编写。在这种情况下，用户可以通过定义一个返回 "real" 的 tracing_mode 方法来禁用伪类化要求。

std::string tracing_mode() {
    return "real";
}

关于伪类化（fakification）的一个注意事项是**张量别名**。我们假设 torchbind 对象内的任何张量都不会与 torchbind 对象外的张量产生别名。因此，改变这些张量中的任何一个都会导致未定义的行为。

在 torch.compile/torch.export 中支持自定义 C++ 类#

为什么我们需要创建一个伪类？#

文档

教程

资源