PyTorch 2.2 中文官方教程(十二)(2)https://developer.aliyun.com/article/1482558
结论
本教程向您展示了如何在 C++中实现自定义 TorchScript 运算符,如何将其构建为共享库,如何在 Python 中使用它来定义 TorchScript 模型,最后如何将其加载到用于推理工作负载的 C++应用程序中。您现在已经准备好通过 C++运算符扩展您的 TorchScript 模型,这些运算符与第三方 C++库进行接口,编写自定义高性能 CUDA 内核,或实现任何其他需要 Python、TorchScript 和 C++之间无缝融合的用例。
如往常一样,如果遇到任何问题或有疑问,您可以使用我们的论坛或GitHub 问题联系我们。此外,我们的常见问题(FAQ)页面可能会提供有用的信息。
附录 A:构建自定义运算符的更多方法
“构建自定义运算符”部分解释了如何使用 CMake 将自定义运算符构建为共享库。本附录概述了两种进一步的编译方法。它们都使用 Python 作为编译过程的“驱动程序”或“接口”。此外,它们都重用了 PyTorch 为C++扩展提供的现有基础设施,这些扩展是依赖于pybind11的 TorchScript 自定义运算符的等效版本,用于将 C++函数“显式”绑定到 Python 中。
第一种方法使用 C++扩展的方便的即时(JIT)编译接口在您首次运行 PyTorch 脚本时在后台编译您的代码。第二种方法依赖于古老的setuptools
包,并涉及编写一个单独的setup.py
文件。这允许更高级的配置以及与其他基于setuptools
的项目集成。我们将在下面详细探讨这两种方法。
使用 JIT 编译进行构建
PyTorch C++扩展工具包提供的 JIT 编译功能允许将自定义运算符的编译直接嵌入到您的 Python 代码中,例如在您的训练脚本顶部。
注意
这里的“JIT 编译”与 TorchScript 编译器中进行的 JIT 编译优化程序无关。它只是意味着您的自定义运算符 C++代码将在您首次导入时编译到系统的/tmp 目录下的一个文件夹中,就好像您之前自己编译过一样。
这个 JIT 编译功能有两种方式。在第一种方式中,您仍然将您的运算符实现放在一个单独的文件中(op.cpp
),然后使用torch.utils.cpp_extension.load()
来编译您的扩展。通常,这个函数会返回暴露您的 C++扩展的 Python 模块。然而,由于我们没有将自定义运算符编译成自己的 Python 模块,我们只想编译一个普通的共享库。幸运的是,torch.utils.cpp_extension.load()
有一个参数is_python_module
,我们可以将其设置为False
,以指示我们只对构建共享库感兴趣,而不是 Python 模块。torch.utils.cpp_extension.load()
然后会编译并加载共享库到当前进程中,就像之前torch.ops.load_library
做的那样:
import torch.utils.cpp_extension torch.utils.cpp_extension.load( name="warp_perspective", sources=["op.cpp"], extra_ldflags=["-lopencv_core", "-lopencv_imgproc"], is_python_module=False, verbose=True ) print(torch.ops.my_ops.warp_perspective)
这应该大致打印:
<built-in method my_ops::warp_perspective of PyCapsule object at 0x7f3e0f840b10>
第二种 JIT 编译的方式允许您将自定义 TorchScript 运算符的源代码作为字符串传递。为此,请使用torch.utils.cpp_extension.load_inline
:
import torch import torch.utils.cpp_extension op_source = """ #include <opencv2/opencv.hpp> #include <torch/script.h> torch::Tensor warp_perspective(torch::Tensor image, torch::Tensor warp) { cv::Mat image_mat(/*rows=*/image.size(0), /*cols=*/image.size(1), /*type=*/CV_32FC1, /*data=*/image.data<float>()); cv::Mat warp_mat(/*rows=*/warp.size(0), /*cols=*/warp.size(1), /*type=*/CV_32FC1, /*data=*/warp.data<float>()); cv::Mat output_mat; cv::warpPerspective(image_mat, output_mat, warp_mat, /*dsize=*/{64, 64}); torch::Tensor output = torch::from_blob(output_mat.ptr<float>(), /*sizes=*/{64, 64}); return output.clone(); } TORCH_LIBRARY(my_ops, m) { m.def("warp_perspective", &warp_perspective); } """ torch.utils.cpp_extension.load_inline( name="warp_perspective", cpp_sources=op_source, extra_ldflags=["-lopencv_core", "-lopencv_imgproc"], is_python_module=False, verbose=True, ) print(torch.ops.my_ops.warp_perspective)
自然地,最佳实践是只在您的源代码相当短的情况下使用torch.utils.cpp_extension.load_inline
。
请注意,如果您在 Jupyter Notebook 中使用这个功能,不要多次执行注册单元格,因为每次执行都会注册一个新的库并重新注册自定义运算符。如果需要重新执行,请在此之前重新启动笔记本的 Python 内核。
使用 Setuptools 构建
从 Python 中构建我们的自定义运算符的第二种方法是使用setuptools
。这样做的好处是setuptools
具有一个非常强大和广泛的接口,用于构建用 C++编写的 Python 模块。然而,由于setuptools
实际上是用于构建 Python 模块而不是普通的共享库(这些库没有模块所需的入口点),这条路线可能有点古怪。也就是说,您只需要一个setup.py
文件来替代CMakeLists.txt
,它看起来像这样:
from setuptools import setup from torch.utils.cpp_extension import BuildExtension, CppExtension setup( name="warp_perspective", ext_modules=[ CppExtension( "warp_perspective", ["example_app/warp_perspective/op.cpp"], libraries=["opencv_core", "opencv_imgproc"], ) ], cmdclass={"build_ext": BuildExtension.with_options(no_python_abi_suffix=True)}, )
请注意,在底部的BuildExtension
中启用了no_python_abi_suffix
选项。这指示setuptools
在生成的共享库名称中省略任何 Python-3 特定的 ABI 后缀。否则,在 Python 3.7 中,库可能被称为warp_perspective.cpython-37m-x86_64-linux-gnu.so
,其中cpython-37m-x86_64-linux-gnu
是 ABI 标签,但我们真的只想让它被称为warp_perspective.so
。
如果我们现在在包含setup.py
的文件夹中的终端中运行python setup.py build develop
,我们应该会看到类似以下的内容:
$ python setup.py build develop running build running build_ext building 'warp_perspective' extension creating build creating build/temp.linux-x86_64-3.7 gcc -pthread -B /root/local/miniconda/compiler_compat -Wl,--sysroot=/ -Wsign-compare -DNDEBUG -g -fwrapv -O3 -Wall -Wstrict-prototypes -fPIC -I/root/local/miniconda/lib/python3.7/site-packages/torch/lib/include -I/root/local/miniconda/lib/python3.7/site-packages/torch/lib/include/torch/csrc/api/include -I/root/local/miniconda/lib/python3.7/site-packages/torch/lib/include/TH -I/root/local/miniconda/lib/python3.7/site-packages/torch/lib/include/THC -I/root/local/miniconda/include/python3.7m -c op.cpp -o build/temp.linux-x86_64-3.7/op.o -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=warp_perspective -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11 cc1plus: warning: command line option ‘-Wstrict-prototypes’ is valid for C/ObjC but not for C++ creating build/lib.linux-x86_64-3.7 g++ -pthread -shared -B /root/local/miniconda/compiler_compat -L/root/local/miniconda/lib -Wl,-rpath=/root/local/miniconda/lib -Wl,--no-as-needed -Wl,--sysroot=/ build/temp.linux-x86_64-3.7/op.o -lopencv_core -lopencv_imgproc -o build/lib.linux-x86_64-3.7/warp_perspective.so running develop running egg_info creating warp_perspective.egg-info writing warp_perspective.egg-info/PKG-INFO writing dependency_links to warp_perspective.egg-info/dependency_links.txt writing top-level names to warp_perspective.egg-info/top_level.txt writing manifest file 'warp_perspective.egg-info/SOURCES.txt' reading manifest file 'warp_perspective.egg-info/SOURCES.txt' writing manifest file 'warp_perspective.egg-info/SOURCES.txt' running build_ext copying build/lib.linux-x86_64-3.7/warp_perspective.so -> Creating /root/local/miniconda/lib/python3.7/site-packages/warp-perspective.egg-link (link to .) Adding warp-perspective 0.0.0 to easy-install.pth file Installed /warp_perspective Processing dependencies for warp-perspective==0.0.0 Finished processing dependencies for warp-perspective==0.0.0
这将生成一个名为warp_perspective.so
的共享库,我们可以像之前那样将其传递给torch.ops.load_library
,以使我们的运算符对 TorchScript 可见:
>>> import torch >>> torch.ops.load_library("warp_perspective.so") >>> print(torch.ops.my_ops.warp_perspective) <built-in method custom::warp_perspective of PyCapsule object at 0x7ff51c5b7bd0>
使用自定义 C++类扩展 TorchScript
原文:
pytorch.org/tutorials/advanced/torch_script_custom_classes.html
译者:飞龙
本教程是自定义运算符教程的后续,介绍了我们为将 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++进程中使用自定义注册的 C++类使用 libtorch。例如,让我们定义一个简单的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
令人难以置信!
将自定义类移动到/从 IValues
还可能需要将自定义类移入或移出IValue
,例如当您从 TorchScript 方法中获取或返回IValue
时,或者您想在 C++中实例化自定义类属性时。要从自定义 C++类实例创建IValue
:
torch::make_custom_class()
提供了类似于 c10::intrusive_ptr的 API,它将接受您提供的一组参数,调用与该参数集匹配的 T 的构造函数,并将该实例包装起来并返回。但是,与仅返回自定义类对象的指针不同,它返回包装对象的IValue
。然后,您可以直接将此IValue
传递给 TorchScript。- 如果您已经有一个指向您的类的
intrusive_ptr
,则可以直接使用构造函数IValue(intrusive_ptr)
从中构造一个 IValue。
将IValue
转换回自定义类:
IValue::toCustomClass()
将返回指向IValue
包含的自定义类的intrusive_ptr
。在内部,此函数正在检查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++类中保存哪些信息。您必须手动指定。方法是在类上使用class_
的特殊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++库进行交互的 C++类来扩展您的 TorchScript 模型,或者实现任何其他需要在 Python、TorchScript 和 C++之间平滑过渡的用例。