本教程是自定義運算符教程的後續教程，並介紹了我們為將 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 Stack : torch::jit::CustomClassHolder {
std::vector<T> stack_;
Stack(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<Stack> clone() const {
return c10::make_intrusive<Stack>(stack_);
}
void merge(const c10::intrusive_ptr<Stack>& c) {
for (auto& elem : c->stack_) {
push(elem);
}
}
};Copy
有幾件事要注意：
torch/custom_class.h是您需要使用自定義類擴展 TorchScript 的標頭。
注意，無論何時使用自定義類的實例，我們都通過c10::intrusive_ptr<>的實例來實現。將intrusive_ptr視為類似於std::shared_ptr的智能指針。使用此智能指針的原因是為了確保在語言(C ++，Python 和 TorchScript）之間對對象實例進行一致的生命周期管理。
注意的第二件事是用戶定義的類必須繼承自torch::jit::CustomClassHolder。這確保了所有設置都可以處理前面提到的生命周期管理系統。
現在讓我們看一下如何使該類對 TorchScript 可見，該過程稱為綁定該類：
// Notice a few things:
// - We pass the class to be registered as a template parameter to
// `torch::jit::class_`. In this instance, we've passed the
// specialization of the Stack class ``Stack<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 single parameter to ``torch::jit::class_()`` is a
// string indicating the name of the class. This is the name
// the class will appear as in both Python and TorchScript.
// For example, our Stack class would appear as ``torch.classes.Stack``.
static auto testStack =
torch::jit::class_<Stack<std::string>>("Stack")
// The following line registers the contructor of our Stack
// class that takes a single `std::vector<std::string>` argument,
// i.e. it exposes the C++ method `Stack(std::vector<T> init)`.
// Currently, we do not support registering overloaded
// constructors, so for now you can only `def()` one instance of
// `torch::jit::init`.
.def(torch::jit::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<Stack<std::string>>& self) {
return self->stack_.back();
})
// The following four lines expose methods of the Stack<std::string>
// class as-is. `torch::jit::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", &Stack<std::string>::push)
.def("pop", &Stack<std::string>::pop)
.def("clone", &Stack<std::string>::clone)
.def("merge", &Stack<std::string>::merge);Copy
現在，我們將使用 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}")Copy
另外，創建一個build目錄。您的文件樹應如下所示：
custom_class_project/
class.cpp
CMakeLists.txt
build/Copy
現在，要構建項目，請繼續從 PyTorch 網站下載適當的 libtorch 二進位文件。將 zip 存檔解壓縮到某個位置(在項目目錄中可能很方便），並記下將其解壓縮到的路徑。接下來，繼續調用 cmake，然後進行構建項目：
$ cd build
$ cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
-- 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_classCopy
您會發現，構建目錄中現在有一個動態庫文件。在 Linux 上，它可能名為libcustom_class.so。因此，文件樹應如下所示：
custom_class_project/
class.cpp
CMakeLists.txt
build/
libcustom_class.soCopy
現在我們已經將我們的類及其註冊編譯為.so文件，我們可以將 <cite>.so</cite> 加載到 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("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 Stack(std::vector<T> init) constructor
# we registered earlier
s = torch.classes.Stack(["foo", "bar"])
# We can call methods in Python
s.push("pushed")
assert s.pop() == "pushed"
# 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.
Stack = torch.classes.Stack
@torch.jit.script
def do_stacks(s : Stack): # We can pass a custom class instance to TorchScript
s2 = torch.classes.Stack(["hi", "mom"]) # We can instantiate the class
s2.merge(s) # We can call a method on the class
return s2.clone(), s2.top() # We can also return instances of the class
# from TorchScript function/methods
stack, top = do_stacks(torch.classes.Stack(["wow"]))
assert top == "wow"
for expected in ["wow", "mom", "hi"]:
assert stack.pop() == expectedCopy
我們也可以在使用 libtorch 的 C ++進程中使用自定義註冊的 C ++類。舉例來說，讓我們定義一個簡單的nn.Module，該實例在我們的 Stack 類上實例化並調用一個方法：
import torch
torch.classes.load_library('libcustom_class.so')
class Foo(torch.nn.Module):
def __init__(self):
super().__init__()
def forward(self, s : str) -> str:
stack = torch.classes.Stack(["hi", "mom"])
return stack.pop() + s
scripted_foo = torch.jit.script(Foo())
print(scripted_foo.graph)
scripted_foo.save('foo.pt')Copy
我們文件系統中的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/Copy
請注意，我們已經複製了序列化的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::script::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;
}Copy
同樣，讓我們定義我們的 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)Copy
您知道練習：cd build，cmake和make：
$ cd build
$ cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
-- 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 inferCopy
現在我們可以運行令人興奮的 C ++二進位文件：
$ ./infer
momfoobarbazCopy
難以置信！
定義自定義 C ++類的序列化/反序列化方法如果您嘗試將具有自定義綁定 C ++類的ScriptModule保存為屬性，則會出現以下錯誤：
# export_attr.py
import torch
torch.classes.load_library('libcustom_class.so')
class Foo(torch.nn.Module):
def __init__(self):
super().__init__()
self.stack = torch.classes.Stack(["just", "testing"])
def forward(self, s : str) -> str:
return self.stack.pop() + s
scripted_foo = torch.jit.script(Foo())
scripted_foo.save('foo.pt')Copy
$ python export_attr.py
RuntimeError: Cannot serialize custom bound C++ class __torch__.torch.classes.Stack. Please define serialization methods via torch::jit::pickle_ for this class. (pushIValueImpl at ../torch/csrc/jit/pickler.cpp:128)Copy
這是因為 TorchScript 無法自動找出 C ++類中保存的信息。您必須手動指定。這樣做的方法是使用class_上的特殊def_pickle方法在類上定義__getstate__和__setstate__方法。
注意
TorchScript 中__getstate__和__setstate__的語義與 Python pickle 模塊的語義相同。您可以閱讀更多有關如何使用這些方法的信息。
這是一個如何更新Stack類的註冊碼以包含序列化方法的示例：
static auto testStack =
torch::jit::class_<Stack<std::string>>("Stack")
.def(torch::jit::init<std::vector<std::string>>())
.def("top", [](const c10::intrusive_ptr<Stack<std::string>>& self) {
return self->stack_.back();
})
.def("push", &Stack<std::string>::push)
.def("pop", &Stack<std::string>::pop)
.def("clone", &Stack<std::string>::clone)
.def("merge", &Stack<std::string>::merge)
// 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<Stack<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<Stack<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 Stack<std::string>
// and call the single-argument std::vector<std::string>
// constructor with the serialized state.
return c10::make_intrusive<Stack<std::string>>(std::move(state));
});Copy
Note
我們採用與 pickle API 中的 pybind11 不同的方法。pybind11 作為傳遞給class_::def()的特殊功能pybind11::pickle()，為此我們有一個單獨的方法def_pickle。這是因為名稱torch::jit::pickle已經被使用，我們不想引起混淆。
以這種方式定義(反）序列化行為後，腳本現在可以成功運行：
import torch
torch.classes.load_library('libcustom_class.so')
class Foo(torch.nn.Module):
def __init__(self):
super().__init__()
self.stack = torch.classes.Stack(["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())Copy
$ python ../export_attr.py
testingCopy
本教程向您介紹了如何向 TorchScript(以及擴展為 Python）公開 C ++類，如何註冊其方法，如何從 Python 和 TorchScript 使用該類以及如何使用該類保存和加載代碼以及運行該代碼。在獨立的 C ++過程中。現在，您可以使用與第三方 C ++庫接口的 C ++類擴展 TorchScript 模型，或實現需要 Python，TorchScript 和 C ++之間的界線才能平滑融合的任何其他用例。