Table of Contents
  • Tutorials >
  • 使用自定义 C++ 类扩展 TorchScript
Shortcuts
advanced/torch_script_custom_classes

使用自定义 C++ 类扩展 TorchScript

Created On: Jan 23, 2020 | Last Updated: Dec 02, 2024 | Last Verified: Nov 05, 2024

警告

TorchScript 不再处于活跃开发状态。

本教程是 自定义操作符 教程的后续内容，介绍了我们构建的 API，用于同时将 C++ 类绑定到 TorchScript 和 Python。这个 API 非常类似于 pybind11，如果你熟悉该系统，许多概念都可以直接迁移过来。

在 C++ 中实现并绑定类

在本教程中，我们将定义一个简单的 C++ 类，该类在成员变量中保持持久状态。

// This header is all you need to do the C++ portions of this
// tutorial
#include <torch/script.h>
// This header is what defines the custom class registration
// behavior specifically. script.h already includes this, but
// we include it here so you know it exists in case you want
// to look at the API or implementation.
#include <torch/custom_class.h>

#include <string>
#include <vector>

template <class T>
struct MyStackClass : torch::CustomClassHolder {
  std::vector<T> stack_;
  MyStackClass(std::vector<T> init) : stack_(init.begin(), init.end()) {}

  void push(T x) {
    stack_.push_back(x);
  }
  T pop() {
    auto val = stack_.back();
    stack_.pop_back();
    return val;
  }

  c10::intrusive_ptr<MyStackClass> clone() const {
    return c10::make_intrusive<MyStackClass>(stack_);
  }

  void merge(const c10::intrusive_ptr<MyStackClass>& c) {
    for (auto& elem : c->stack_) {
      push(elem);
    }
  }
};

需要注意以下几点：

  • torch/custom_class.h 是扩展 TorchScript 所需包含的头文件。

  • 注意，每当我们使用自定义类的实例时，我们都是通过 c10::intrusive_ptr<> 实例来操作。可将 intrusive_ptr 类比作 std::shared_ptr 的智能指针，但它的引用计数直接存储在对象中，而不是像 std::shared_ptr 那样存储在一个单独的元数据块中。 torch::Tensor 内部使用同样的指针类型；而自定义类也必须使用此指针类型，以便我们能够一致地管理不同对象类型。

  • 第二点需要注意的是用户定义的类必须继承自 torch::CustomClassHolder。这样可以确保自定义类有空间存储引用计数。

现在让我们看看如何使此类对 TorchScript 可见，这一过程称为 绑定 类：

// Notice a few things:
// - We pass the class to be registered as a template parameter to
//   `torch::class_`. In this instance, we've passed the
//   specialization of the MyStackClass class ``MyStackClass<std::string>``.
//   In general, you cannot register a non-specialized template
//   class. For non-templated classes, you can just pass the
//   class name directly as the template parameter.
// - The arguments passed to the constructor make up the "qualified name"
//   of the class. In this case, the registered class will appear in
//   Python and C++ as `torch.classes.my_classes.MyStackClass`. We call
//   the first argument the "namespace" and the second argument the
//   actual class name.
TORCH_LIBRARY(my_classes, m) {
  m.class_<MyStackClass<std::string>>("MyStackClass")
    // The following line registers the contructor of our MyStackClass
    // class that takes a single `std::vector<std::string>` argument,
    // i.e. it exposes the C++ method `MyStackClass(std::vector<T> init)`.
    // Currently, we do not support registering overloaded
    // constructors, so for now you can only `def()` one instance of
    // `torch::init`.
    .def(torch::init<std::vector<std::string>>())
    // The next line registers a stateless (i.e. no captures) C++ lambda
    // function as a method. Note that a lambda function must take a
    // `c10::intrusive_ptr<YourClass>` (or some const/ref version of that)
    // as the first argument. Other arguments can be whatever you want.
    .def("top", [](const c10::intrusive_ptr<MyStackClass<std::string>>& self) {
      return self->stack_.back();
    })
    // The following four lines expose methods of the MyStackClass<std::string>
    // class as-is. `torch::class_` will automatically examine the
    // argument and return types of the passed-in method pointers and
    // expose these to Python and TorchScript accordingly. Finally, notice
    // that we must take the *address* of the fully-qualified method name,
    // i.e. use the unary `&` operator, due to C++ typing rules.
    .def("push", &MyStackClass<std::string>::push)
    .def("pop", &MyStackClass<std::string>::pop)
    .def("clone", &MyStackClass<std::string>::clone)
    .def("merge", &MyStackClass<std::string>::merge)
  ;
}

使用 CMake 将示例构建为 C++ 项目

现在，我们将使用 CMake 构建系统来构建上述 C++代码。首先，将迄今为止介绍的所有 C++代码放入一个名为 class.cpp 的文件中。然后，编写一个简单的 CMakeLists.txt 文件，并将其放在相同目录下。以下是 CMakeLists.txt 的内容：

cmake_minimum_required(VERSION 3.1 FATAL_ERROR)
project(custom_class)

find_package(Torch REQUIRED)

# Define our library target
add_library(custom_class SHARED class.cpp)
set(CMAKE_CXX_STANDARD 14)
# Link against LibTorch
target_link_libraries(custom_class "${TORCH_LIBRARIES}")

此外，创建一个 build 目录。你的文件树应该像这样:

custom_class_project/
  class.cpp
  CMakeLists.txt
  build/

我们假设你按照 之前的教程 设置了你的环境。继续调用 cmake，然后使用 make 来构建项目：

$ cd build
$ cmake -DCMAKE_PREFIX_PATH="$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')" ..
  -- The C compiler identification is GNU 7.3.1
  -- The CXX compiler identification is GNU 7.3.1
  -- Check for working C compiler: /opt/rh/devtoolset-7/root/usr/bin/cc
  -- Check for working C compiler: /opt/rh/devtoolset-7/root/usr/bin/cc -- works
  -- Detecting C compiler ABI info
  -- Detecting C compiler ABI info - done
  -- Detecting C compile features
  -- Detecting C compile features - done
  -- Check for working CXX compiler: /opt/rh/devtoolset-7/root/usr/bin/c++
  -- Check for working CXX compiler: /opt/rh/devtoolset-7/root/usr/bin/c++ -- works
  -- Detecting CXX compiler ABI info
  -- Detecting CXX compiler ABI info - done
  -- Detecting CXX compile features
  -- Detecting CXX compile features - done
  -- Looking for pthread.h
  -- Looking for pthread.h - found
  -- Looking for pthread_create
  -- Looking for pthread_create - not found
  -- Looking for pthread_create in pthreads
  -- Looking for pthread_create in pthreads - not found
  -- Looking for pthread_create in pthread
  -- Looking for pthread_create in pthread - found
  -- Found Threads: TRUE
  -- Found torch: /torchbind_tutorial/libtorch/lib/libtorch.so
  -- Configuring done
  -- Generating done
  -- Build files have been written to: /torchbind_tutorial/build
$ make -j
  Scanning dependencies of target custom_class
  [ 50%] Building CXX object CMakeFiles/custom_class.dir/class.cpp.o
  [100%] Linking CXX shared library libcustom_class.so
  [100%] Built target custom_class

你会发现现在（以及其他文件）在 build 目录下出现了一个动态库文件。在 Linux 上，这个文件很可能名为 libcustom_class.so。因此文件树应该看起来像这样:

custom_class_project/
  class.cpp
  CMakeLists.txt
  build/
    libcustom_class.so

在 Python 和 TorchScript 中使用 C++ 类

现在我们已经将类及其注册编译成一个 .so 文件，我们可以将此 .so 文件加载到 Python 中并试验它。以下是演示代码：

import torch

# `torch.classes.load_library()` allows you to pass the path to your .so file
# to load it in and make the custom C++ classes available to both Python and
# TorchScript
torch.classes.load_library("build/libcustom_class.so")
# You can query the loaded libraries like this:
print(torch.classes.loaded_libraries)
# prints {'/custom_class_project/build/libcustom_class.so'}

# We can find and instantiate our custom C++ class in python by using the
# `torch.classes` namespace:
#
# This instantiation will invoke the MyStackClass(std::vector<T> init)
# constructor we registered earlier
s = torch.classes.my_classes.MyStackClass(["foo", "bar"])

# We can call methods in Python
s.push("pushed")
assert s.pop() == "pushed"

# Test custom operator
s.push("pushed")
torch.ops.my_classes.manipulate_instance(s)  # acting as s.pop()
assert s.top() == "bar" 

# Returning and passing instances of custom classes works as you'd expect
s2 = s.clone()
s.merge(s2)
for expected in ["bar", "foo", "bar", "foo"]:
    assert s.pop() == expected

# We can also use the class in TorchScript
# For now, we need to assign the class's type to a local in order to
# annotate the type on the TorchScript function. This may change
# in the future.
MyStackClass = torch.classes.my_classes.MyStackClass


@torch.jit.script
def do_stacks(s: MyStackClass):  # We can pass a custom class instance
    # We can instantiate the class
    s2 = torch.classes.my_classes.MyStackClass(["hi", "mom"])
    s2.merge(s)  # We can call a method on the class
    # We can also return instances of the class
    # from TorchScript function/methods
    return s2.clone(), s2.top()


stack, top = do_stacks(torch.classes.my_classes.MyStackClass(["wow"]))
assert top == "wow"
for expected in ["wow", "mom", "hi"]:
    assert stack.pop() == expected

使用自定义类保存、加载和运行 TorchScript代码

我们还可以在使用 libtorch 的 C++ 进程中使用自定义注册的 C++类。作为示例，让我们定义一个简单的``nn.Module``，它实例化并调用我们 MyStackClass 类中的方法：

import torch

torch.classes.load_library('build/libcustom_class.so')


class Foo(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, s: str) -> str:
        stack = torch.classes.my_classes.MyStackClass(["hi", "mom"])
        return stack.pop() + s


scripted_foo = torch.jit.script(Foo())
print(scripted_foo.graph)

scripted_foo.save('foo.pt')

现在我们的文件系统中已经包含了我们刚刚定义的序列化 TorchScript 程序 foo.pt。

现在，我们将定义一个新的 CMake 项目，以展示如何加载这个模型及其所需的 .so 文件。关于如何全面实现这一点，请参考 在 C++ 中加载 TorchScript 模型的教程。

与之前类似，让我们创建一个包含以下内容的文件结构:

cpp_inference_example/
  infer.cpp
  CMakeLists.txt
  foo.pt
  build/
  custom_class_project/
    class.cpp
    CMakeLists.txt
    build/

注意我们已经复制了序列化的 foo.pt 文件，以及上述 custom_class_project 的源代码树。我们将添加 custom_class_project 作为此C++ 项目的依赖项，以便将自定义类构建到二进制文件中。

让我们用以下内容填充 infer.cpp：

#include <torch/script.h>

#include <iostream>
#include <memory>

int main(int argc, const char* argv[]) {
  torch::jit::Module module;
  try {
    // Deserialize the ScriptModule from a file using torch::jit::load().
    module = torch::jit::load("foo.pt");
  }
  catch (const c10::Error& e) {
    std::cerr << "error loading the model\n";
    return -1;
  }

  std::vector<c10::IValue> inputs = {"foobarbaz"};
  auto output = module.forward(inputs).toString();
  std::cout << output->string() << std::endl;
}

类似地让我们定义我们的 CMakeLists.txt 文件：

cmake_minimum_required(VERSION 3.1 FATAL_ERROR)
project(infer)

find_package(Torch REQUIRED)

add_subdirectory(custom_class_project)

# Define our library target
add_executable(infer infer.cpp)
set(CMAKE_CXX_STANDARD 14)
# Link against LibTorch
target_link_libraries(infer "${TORCH_LIBRARIES}")
# This is where we link in our libcustom_class code, making our
# custom class available in our binary.
target_link_libraries(infer -Wl,--no-as-needed custom_class)

按照流程： cd build，cmake，然后运行 make：

$ cd build
$ cmake -DCMAKE_PREFIX_PATH="$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')" ..
  -- The C compiler identification is GNU 7.3.1
  -- The CXX compiler identification is GNU 7.3.1
  -- Check for working C compiler: /opt/rh/devtoolset-7/root/usr/bin/cc
  -- Check for working C compiler: /opt/rh/devtoolset-7/root/usr/bin/cc -- works
  -- Detecting C compiler ABI info
  -- Detecting C compiler ABI info - done
  -- Detecting C compile features
  -- Detecting C compile features - done
  -- Check for working CXX compiler: /opt/rh/devtoolset-7/root/usr/bin/c++
  -- Check for working CXX compiler: /opt/rh/devtoolset-7/root/usr/bin/c++ -- works
  -- Detecting CXX compiler ABI info
  -- Detecting CXX compiler ABI info - done
  -- Detecting CXX compile features
  -- Detecting CXX compile features - done
  -- Looking for pthread.h
  -- Looking for pthread.h - found
  -- Looking for pthread_create
  -- Looking for pthread_create - not found
  -- Looking for pthread_create in pthreads
  -- Looking for pthread_create in pthreads - not found
  -- Looking for pthread_create in pthread
  -- Looking for pthread_create in pthread - found
  -- Found Threads: TRUE
  -- Found torch: /local/miniconda3/lib/python3.7/site-packages/torch/lib/libtorch.so
  -- Configuring done
  -- Generating done
  -- Build files have been written to: /cpp_inference_example/build
$ make -j
  Scanning dependencies of target custom_class
  [ 25%] Building CXX object custom_class_project/CMakeFiles/custom_class.dir/class.cpp.o
  [ 50%] Linking CXX shared library libcustom_class.so
  [ 50%] Built target custom_class
  Scanning dependencies of target infer
  [ 75%] Building CXX object CMakeFiles/infer.dir/infer.cpp.o
  [100%] Linking CXX executable infer
  [100%] Built target infer

现在我们可以运行令人激动的 C++ 二进制文件了：

$ ./infer
  momfoobarbaz

不可思议！

将自定义类移动到/从 IValues 中

你可能还需要将自定义类移动到或移出 IValue 中，例如当你从 TorchScript 方法中获取或返回 IValue，或者希望在 C++ 中实例化一个自定义类属性时。对于从自定义 C++ 类实例创建 IValue：

  • torch::make_custom_class<T>() 提供了一个类似于 c10::intrusive_ptr<T> 的 API，它会接收你提供给它的任意参数集，调用匹配该参数集的 T 的构造函数，并包装该实例后返回。然而，它不是仅返回指向自定义类对象的指针，而是返回一个带有对象包装的 IValue。然后你可以直接将这个 IValue 传递给 TorchScript。

  • 如果你已经有一个指向你类的 intrusive_ptr，你可以使用构造函数 IValue(intrusive_ptr<T>) 直接从它构造一个 IValue。

对于将 IValue 转换回自定义类：

  • IValue::toCustomClass<T>() 会返回一个 intrusive_ptr<T>，指向 IValue 包含的自定义类。内部，这个函数会检查 T 是否已经注册为自定义类，并检查 IValue 是否确实包含一个自定义类。你可以通过调用 isCustomClass() 手动检查 IValue 是否包含自定义类。

为自定义 C++ 类定义序列化/反序列化方法

如果你尝试保存一个以自定义绑定 C++ 类为属性的 ScriptModule，你会收到如下错误消息：

# export_attr.py
import torch

torch.classes.load_library('build/libcustom_class.so')


class Foo(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.stack = torch.classes.my_classes.MyStackClass(["just", "testing"])

    def forward(self, s: str) -> str:
        return self.stack.pop() + s


scripted_foo = torch.jit.script(Foo())

scripted_foo.save('foo.pt')
loaded = torch.jit.load('foo.pt')

print(loaded.stack.pop())

$ python export_attr.py
RuntimeError: Cannot serialize custom bound C++ class __torch__.torch.classes.my_classes.MyStackClass. Please define serialization methods via def_pickle for this class. (pushIValueImpl at ../torch/csrc/jit/pickler.cpp:128)

这是因为 TorchScript 无法自动确定要从你的 C++ 类中保存哪些信息。你必须手动指定。实现方式是为类定义 __getstate__ 和 __setstate__ 方法，使用 class_ 的特殊 def_pickle 方法。

备注

TorchScript 中 __getstate__ 和 __setstate__ 的语义与 Python pickle 模块的语义相同。关于我们如何使用这些方法，可以 阅读更多。

以下是我们可以添加到 MyStackClass 注册中的 def_pickle 调用示例，以包含序列化方法：

    // class_<>::def_pickle allows you to define the serialization
    // and deserialization methods for your C++ class.
    // Currently, we only support passing stateless lambda functions
    // as arguments to def_pickle
    .def_pickle(
          // __getstate__
          // This function defines what data structure should be produced
          // when we serialize an instance of this class. The function
          // must take a single `self` argument, which is an intrusive_ptr
          // to the instance of the object. The function can return
          // any type that is supported as a return value of the TorchScript
          // custom operator API. In this instance, we've chosen to return
          // a std::vector<std::string> as the salient data to preserve
          // from the class.
          [](const c10::intrusive_ptr<MyStackClass<std::string>>& self)
              -> std::vector<std::string> {
            return self->stack_;
          },
          // __setstate__
          // This function defines how to create a new instance of the C++
          // class when we are deserializing. The function must take a
          // single argument of the same type as the return value of
          // `__getstate__`. The function must return an intrusive_ptr
          // to a new instance of the C++ class, initialized however
          // you would like given the serialized state.
          [](std::vector<std::string> state)
              -> c10::intrusive_ptr<MyStackClass<std::string>> {
            // A convenient way to instantiate an object and get an
            // intrusive_ptr to it is via `make_intrusive`. We use
            // that here to allocate an instance of MyStackClass<std::string>
            // and call the single-argument std::vector<std::string>
            // constructor with the serialized state.
            return c10::make_intrusive<MyStackClass<std::string>>(std::move(state));
          });

备注

我们在 pickle API 上采用了不同于 pybind11 的方法。pybind11 提供了一个名为 pybind11::pickle() 的特殊函数，你可以将其传递给 class_::def()；而我们为了这个目的有一个单独的方法 def_pickle。这是因为``torch::jit::pickle`` 这个名称已经被占用了，我们不希望引起混乱。

一旦以这种方式定义了（反）序列化行为，我们的脚本现在可以成功运行：

$ python ../export_attr.py
testing

定义采用或返回绑定 C++ 类的自定义操作符

一旦你定义了一个自定义 C++ 类，你也可以使用该类作为参数或返回值来自定义操作符（即自由函数）。假设你有如下自由函数：

c10::intrusive_ptr<MyStackClass<std::string>> manipulate_instance(const c10::intrusive_ptr<MyStackClass<std::string>>& instance) {
  instance->pop();
  return instance;
}

你可以在你的 TORCH_LIBRARY 块内运行以下代码进行注册：

    m.def(
      "manipulate_instance(__torch__.torch.classes.my_classes.MyStackClass x) -> __torch__.torch.classes.my_classes.MyStackClass Y",
      manipulate_instance
    );

有关注册 API 的详细信息，请参考 自定义操作符教程。

完成后，你可以像以下示例一样使用该操作符：

class TryCustomOp(torch.nn.Module):
    def __init__(self):
        super(TryCustomOp, self).__init__()
        self.f = torch.classes.my_classes.MyStackClass(["foo", "bar"])

    def forward(self):
        return torch.ops.my_classes.manipulate_instance(self.f)

备注

注册一个以 C++ 类为参数的操作符要求该自定义类已被注册。你可以通过确保自定义类注册与你的自由函数定义在同一个 TORCH_LIBRARY 块中，并且自定义类注册优先来强制执行这一点。在未来，我们可能会放松这个要求，以允许这些可以按任何顺序注册。

总结

本教程向你介绍了如何将 C++ 类公开给 TorchScript（以及扩展到 Python）、如何注册其方法、如何从 Python 和 TorchScript 使用该类，以及如何使用该类保存和加载代码，并在独立的 C++ 进程中运行该代码。你现在已经准备好通过接口第三方 C++库或实现任何其他需要在 Python、TorchScript 和 C++之间平滑交互的用例，使用 C++ 类扩展你的 TorchScript 模型了。

一如既往，如果您遇到任何问题或有任何疑问，可以使用我们的 论坛 或 GitHub问题 联系我们。此外，我们的 常见问题解答 (FAQ) 页 可能也有有用的信息。

文档

访问 PyTorch 的详细开发者文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源

为了分析流量并优化您的体验，我们在此网站上提供Cookie。点击或浏览即表示您同意我们使用Cookie。 作为本网站的当前维护者，Facebook的Cookie政策适用。了解更多信息，包括可用的控件： Cookies政策.