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

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

警告

此功能处于原型状态，并且可能发生向后兼容性破坏性更改。本教程提供了 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)

为什么要创建虚类？#

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

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

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

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

虚化的一个注意事项是关于**张量别名**。我们假设 torchbind 对象中的张量不会别名 torchbind 对象之外的张量。因此，修改其中一个张量将导致未定义的行为。

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

为什么要创建虚类？#

文档

教程

资源