目錄
  • 教程 >
  • 使用自定義 C++ 類別擴展 TorchScript
捷徑
advanced/torch_script_custom_classes

使用自訂 C++ 類別擴展 TorchScript

建立於：2020 年 1 月 23 日 | 最後更新：2024 年 12 月 02 日 | 最後驗證：2024 年 11 月 05 日

警告

TorchScript 已停止積極開發。

本教程是自訂運算子教程的後續內容，介紹了我們為同時將 C++ 類別綁定到 TorchScript 和 Python 中而建立的 API。 此 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，然後進行建置專案

$ 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

您會發現現在（除其他事項外）建置目錄中存在一個動態程式庫檔案。 在 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 程式碼

我們還可以在 C++ 流程中使用 libtorch 中的自訂註冊 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')

現在，我們的檔案系統中的 foo.pt 包含我們剛剛定義的序列化 TorchScript 程式。

現在，我們將定義一個新的 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

太棒了！

自定義類別與 IValue 之間的移動

您可能也需要將自定義類別移入或移出 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>() 將返回一個指向 IValue 包含的自定義類別的 intrusive_ptr<T>。在內部，此函數會檢查 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++ 類別中儲存哪些資訊。您必須手動指定。方法是在類別上使用特殊的 def_pickle 方法定義 __getstate__ 和 __setstate__ 方法。

注意

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++ 類別擴展您的 TorchScript 模型，該類別與第三方 C++ 函式庫介接，或實作任何其他需要 Python、TorchScript 和 C++ 之間的界線順暢融合的用例。

與往常一樣，如果您遇到任何問題或有疑問，您可以使用我們的 論壇 或 GitHub 問題 與我們聯繫。此外，我們的 常見問題解答 (FAQ) 頁面 可能會有有用的資訊。

文件

存取 PyTorch 的完整開發人員文件

檢視文件

教程

取得初學者和進階開發人員的深入教學課程

檢視教學課程

資源

尋找開發資源並取得問題解答

檢視資源

為了分析流量並優化您的體驗，我們在本網站上使用 Cookie。點擊或瀏覽即表示您同意允許我們使用 Cookie。作為本網站目前的維護者，Facebook 的 Cookie 政策適用。了解更多，包括可用的控制項：Cookie 政策。