PyTorch 使用自定義 C ++類擴展 TorchScript

2020-09-10 10:40 更新
原文： https://pytorch.org/tutorials/advanced/torch_script_custom_classes.html

本教程是自定義運算符教程的后續(xù)教程，并介紹了我們?yōu)閷?C ++類同時綁定到 TorchScript 和 Python 而構(gòu)建的 API。 該 API 與 pybind11 非常相似，如果您熟悉該系統(tǒng)，則大多數(shù)概念都將轉(zhuǎn)移過來。

在 C ++中實現(xiàn)和綁定類

在本教程中，我們將定義一個簡單的 C ++類，該類在成員變量中保持持久狀態(tài)。

  1. // This header is all you need to do the C++ portions of this
  2. // tutorial
  3. #include <torch/script.h>
  4. // This header is what defines the custom class registration
  5. // behavior specifically. script.h already includes this, but
  6. // we include it here so you know it exists in case you want
  7. // to look at the API or implementation.
  8. #include <torch/custom_class.h>
  11. #include <string>
  12. #include <vector>
  15. template <class T>
  16. struct Stack : torch::jit::CustomClassHolder {
  17.   std::vector<T> stack_;
  18.   Stack(std::vector<T> init) : stack_(init.begin(), init.end()) {}
  21.   void push(T x) {
  22.     stack_.push_back(x);
  23.   }
  24.   T pop() {
  25.     auto val = stack_.back();
  26.     stack_.pop_back();
  27.     return val;
  28.   }
  31.   c10::intrusive_ptr<Stack> clone() const {
  32.     return c10::make_intrusive<Stack>(stack_);
  33.   }
  36.   void merge(const c10::intrusive_ptr<Stack>& c) {
  37.     for (auto& elem : c->stack_) {
  38.       push(elem);
  39.     }
  40.   }
  41. };

有幾件事要注意：

  • torch/custom_class.h是您需要使用自定義類擴展 TorchScript 的標(biāo)頭。
  • 注意，無論何時使用自定義類的實例，我們都通過c10::intrusive_ptr&lt;&gt;的實例來實現(xiàn)。 將intrusive_ptr視為類似于std::shared_ptr的智能指針。 使用此智能指針的原因是為了確保在語言(C ++，Python 和 TorchScript）之間對對象實例進(jìn)行一致的生命周期管理。
  • 注意的第二件事是用戶定義的類必須繼承自torch::jit::CustomClassHolder。 這確保了所有設(shè)置都可以處理前面提到的生命周期管理系統(tǒng)。

現(xiàn)在讓我們看一下如何使該類對 TorchScript 可見，該過程稱為綁定該類：

  1. // Notice a few things:
  2. // - We pass the class to be registered as a template parameter to
  3. //   `torch::jit::class_`. In this instance, we've passed the
  4. //   specialization of the Stack class ``Stack<std::string>``.
  5. //   In general, you cannot register a non-specialized template
  6. //   class. For non-templated classes, you can just pass the
  7. //   class name directly as the template parameter.
  8. // - The single parameter to ``torch::jit::class_()`` is a
  9. //   string indicating the name of the class. This is the name
  10. //   the class will appear as in both Python and TorchScript.
  11. //   For example, our Stack class would appear as ``torch.classes.Stack``.
  12. static auto testStack =
  13.   torch::jit::class_<Stack<std::string>>("Stack")
  14.       // The following line registers the contructor of our Stack
  15.       // class that takes a single `std::vector<std::string>` argument,
  16.       // i.e. it exposes the C++ method `Stack(std::vector<T> init)`.
  17.       // Currently, we do not support registering overloaded
  18.       // constructors, so for now you can only `def()` one instance of
  19.       // `torch::jit::init`.
  20.       .def(torch::jit::init<std::vector<std::string>>())
  21.       // The next line registers a stateless (i.e. no captures) C++ lambda
  22.       // function as a method. Note that a lambda function must take a
  23.       // `c10::intrusive_ptr<YourClass>` (or some const/ref version of that)
  24.       // as the first argument. Other arguments can be whatever you want.
  25.       .def("top", [](const c10::intrusive_ptr<Stack<std::string>>& self) {
  26.         return self->stack_.back();
  27.       })
  28.       // The following four lines expose methods of the Stack<std::string>
  29.       // class as-is. `torch::jit::class_` will automatically examine the
  30.       // argument and return types of the passed-in method pointers and
  31.       // expose these to Python and TorchScript accordingly. Finally, notice
  32.       // that we must take the *address* of the fully-qualified method name,
  33.       // i.e. use the unary `&` operator, due to C++ typing rules.
  34.       .def("push", &Stack<std::string>::push)
  35.       .def("pop", &Stack<std::string>::pop)
  36.       .def("clone", &Stack<std::string>::clone)
  37.       .def("merge", &Stack<std::string>::merge);

使用 CMake 將示例構(gòu)建為 C ++項目

現(xiàn)在，我們將使用 CMake 構(gòu)建系統(tǒng)來構(gòu)建上述 C ++代碼。 首先，將到目前為止介紹的所有 C ++代碼放入class.cpp/文件中。 然后，編寫一個簡單的CMakeLists.txt文件并將其放置在同一目錄中。 CMakeLists.txt的外觀如下： 

  1. cmake_minimum_required(VERSION 3.1 FATAL_ERROR)
  2. project(custom_class)
  5. find_package(Torch REQUIRED)
  8. ## Define our library target
  9. add_library(custom_class SHARED class.cpp)
  10. set(CMAKE_CXX_STANDARD 14)
  11. ## Link against LibTorch
  12. target_link_libraries(custom_class "${TORCH_LIBRARIES}")

另外，創(chuàng)建一個build目錄。 您的文件樹應(yīng)如下所示：

  1. custom_class_project/
  2.   class.cpp
  3.   CMakeLists.txt
  4.   build/

現(xiàn)在，要構(gòu)建項目，請繼續(xù)從 PyTorch 網(wǎng)站下載適當(dāng)?shù)?libtorch 二進(jìn)制文件。 將 zip 存檔解壓縮到某個位置(在項目目錄中可能很方便），并記下將其解壓縮到的路徑。 接下來，繼續(xù)調(diào)用 cmake，然后進(jìn)行構(gòu)建項目：

  1. $ cd build
  2. $ cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
  3.   -- The C compiler identification is GNU 7.3.1
  4.   -- The CXX compiler identification is GNU 7.3.1
  5.   -- Check for working C compiler: /opt/rh/devtoolset-7/root/usr/bin/cc
  6.   -- Check for working C compiler: /opt/rh/devtoolset-7/root/usr/bin/cc -- works
  7.   -- Detecting C compiler ABI info
  8.   -- Detecting C compiler ABI info - done
  9.   -- Detecting C compile features
  10.   -- Detecting C compile features - done
  11.   -- Check for working CXX compiler: /opt/rh/devtoolset-7/root/usr/bin/c++
  12.   -- Check for working CXX compiler: /opt/rh/devtoolset-7/root/usr/bin/c++ -- works
  13.   -- Detecting CXX compiler ABI info
  14.   -- Detecting CXX compiler ABI info - done
  15.   -- Detecting CXX compile features
  16.   -- Detecting CXX compile features - done
  17.   -- Looking for pthread.h
  18.   -- Looking for pthread.h - found
  19.   -- Looking for pthread_create
  20.   -- Looking for pthread_create - not found
  21.   -- Looking for pthread_create in pthreads
  22.   -- Looking for pthread_create in pthreads - not found
  23.   -- Looking for pthread_create in pthread
  24.   -- Looking for pthread_create in pthread - found
  25.   -- Found Threads: TRUE
  26.   -- Found torch: /torchbind_tutorial/libtorch/lib/libtorch.so
  27.   -- Configuring done
  28.   -- Generating done
  29.   -- Build files have been written to: /torchbind_tutorial/build
  30. $ make -j
  31.   Scanning dependencies of target custom_class
  32.   [ 50%] Building CXX object CMakeFiles/custom_class.dir/class.cpp.o
  33.   [100%] Linking CXX shared library libcustom_class.so
  34.   [100%] Built target custom_class

您會發(fā)現(xiàn)，構(gòu)建目錄中現(xiàn)在有一個動態(tài)庫文件。 在 Linux 上，它可能名為libcustom_class.so。 因此，文件樹應(yīng)如下所示：

  1. custom_class_project/
  2.   class.cpp
  3.   CMakeLists.txt
  4.   build/
  5.     libcustom_class.so

從 Python 和 TorchScript 使用 C ++類

現(xiàn)在我們已經(jīng)將我們的類及其注冊編譯為.so文件，我們可以將 <cite>.so</cite> 加載到 Python 中并進(jìn)行嘗試。 這是一個演示腳本的腳本：

  1. import torch
  4. ## `torch.classes.load_library()` allows you to pass the path to your .so file
  5. ## to load it in and make the custom C++ classes available to both Python and
  6. ## TorchScript
  7. torch.classes.load_library("libcustom_class.so")
  8. ## You can query the loaded libraries like this:
  9. print(torch.classes.loaded_libraries)
  10. ## prints {'/custom_class_project/build/libcustom_class.so'}
  13. ## We can find and instantiate our custom C++ class in python by using the
  14. ## `torch.classes` namespace:
  15. ## ## This instantiation will invoke the Stack(std::vector<T> init) constructor
  16. ## we registered earlier
  17. s = torch.classes.Stack(["foo", "bar"])
  20. ## We can call methods in Python
  21. s.push("pushed")
  22. assert s.pop() == "pushed"
  25. ## Returning and passing instances of custom classes works as you'd expect
  26. s2 = s.clone()
  27. s.merge(s2)
  28. for expected in ["bar", "foo", "bar", "foo"]:
  29.     assert s.pop() == expected
  32. ## We can also use the class in TorchScript
  33. ## For now, we need to assign the class's type to a local in order to
  34. ## annotate the type on the TorchScript function. This may change
  35. ## in the future.
  36. Stack = torch.classes.Stack
  39. @torch.jit.script
  40. def do_stacks(s : Stack): # We can pass a custom class instance to TorchScript
  41.     s2 = torch.classes.Stack(["hi", "mom"]) # We can instantiate the class
  42.     s2.merge(s) # We can call a method on the class
  43.     return s2.clone(), s2.top()  # We can also return instances of the class
  44.                                  # from TorchScript function/methods
  47. stack, top = do_stacks(torch.classes.Stack(["wow"]))
  48. assert top == "wow"
  49. for expected in ["wow", "mom", "hi"]:
  50.     assert stack.pop() == expected

使用自定義類保存，加載和運行 TorchScript 代碼

我們也可以在使用 libtorch 的 C ++進(jìn)程中使用自定義注冊的 C ++類。 舉例來說，讓我們定義一個簡單的nn.Module，該實例在我們的 Stack 類上實例化并調(diào)用一個方法：

  1. import torch
  4. torch.classes.load_library('libcustom_class.so')
  7. class Foo(torch.nn.Module):
  8.     def __init__(self):
  9.         super().__init__()
  12.     def forward(self, s : str) -> str:
  13.         stack = torch.classes.Stack(["hi", "mom"])
  14.         return stack.pop() + s
  17. scripted_foo = torch.jit.script(Foo())
  18. print(scripted_foo.graph)
  21. scripted_foo.save('foo.pt')

我們文件系統(tǒng)中的foo.pt現(xiàn)在包含我們剛剛定義的序列化 TorchScript 程序。

現(xiàn)在，我們將定義一個新的 CMake 項目，以展示如何加載此模型及其所需的.so 文件。 有關(guān)如何執(zhí)行此操作的完整說明，請查看在 C ++教程中加載 TorchScript 模型。

與之前類似，讓我們創(chuàng)建一個包含以下內(nèi)容的文件結(jié)構(gòu)：

  1. cpp_inference_example/
  2.   infer.cpp
  3.   CMakeLists.txt
  4.   foo.pt
  5.   build/
  6.   custom_class_project/
  7.     class.cpp
  8.     CMakeLists.txt
  9.     build/

請注意，我們已經(jīng)復(fù)制了序列化的foo.pt文件以及上面custom_class_project的源代碼樹。 我們將添加custom_class_project作為對此 C ++項目的依賴項，以便我們可以將自定義類構(gòu)建到二進(jìn)制文件中。

讓我們用以下內(nèi)容填充infer.cpp：

  1. #include <torch/script.h>
  4. #include <iostream>
  5. #include <memory>
  8. int main(int argc, const char* argv[]) {
  9.   torch::jit::script::Module module;
  10.   try {
  11.     // Deserialize the ScriptModule from a file using torch::jit::load().
  12.     module = torch::jit::load("foo.pt");
  13.   }
  14.   catch (const c10::Error& e) {
  15.     std::cerr << "error loading the model\n";
  16.     return -1;
  17.   }
  20.   std::vector<c10::IValue> inputs = {"foobarbaz"};
  21.   auto output = module.forward(inputs).toString();
  22.   std::cout << output->string() << std::endl;
  23. }

同樣，讓我們??定義我們的 CMakeLists.txt 文件：

  1. cmake_minimum_required(VERSION 3.1 FATAL_ERROR)
  2. project(infer)
  5. find_package(Torch REQUIRED)
  8. add_subdirectory(custom_class_project)
  11. ## Define our library target
  12. add_executable(infer infer.cpp)
  13. set(CMAKE_CXX_STANDARD 14)
  14. ## Link against LibTorch
  15. target_link_libraries(infer "${TORCH_LIBRARIES}")
  16. ## This is where we link in our libcustom_class code, making our
  17. ## custom class available in our binary.
  18. target_link_libraries(infer -Wl,--no-as-needed custom_class)

您知道練習(xí)：cd build，cmake和make：

  1. $ cd build
  2. $ cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
  3.   -- The C compiler identification is GNU 7.3.1
  4.   -- The CXX compiler identification is GNU 7.3.1
  5.   -- Check for working C compiler: /opt/rh/devtoolset-7/root/usr/bin/cc
  6.   -- Check for working C compiler: /opt/rh/devtoolset-7/root/usr/bin/cc -- works
  7.   -- Detecting C compiler ABI info
  8.   -- Detecting C compiler ABI info - done
  9.   -- Detecting C compile features
  10.   -- Detecting C compile features - done
  11.   -- Check for working CXX compiler: /opt/rh/devtoolset-7/root/usr/bin/c++
  12.   -- Check for working CXX compiler: /opt/rh/devtoolset-7/root/usr/bin/c++ -- works
  13.   -- Detecting CXX compiler ABI info
  14.   -- Detecting CXX compiler ABI info - done
  15.   -- Detecting CXX compile features
  16.   -- Detecting CXX compile features - done
  17.   -- Looking for pthread.h
  18.   -- Looking for pthread.h - found
  19.   -- Looking for pthread_create
  20.   -- Looking for pthread_create - not found
  21.   -- Looking for pthread_create in pthreads
  22.   -- Looking for pthread_create in pthreads - not found
  23.   -- Looking for pthread_create in pthread
  24.   -- Looking for pthread_create in pthread - found
  25.   -- Found Threads: TRUE
  26.   -- Found torch: /local/miniconda3/lib/python3.7/site-packages/torch/lib/libtorch.so
  27.   -- Configuring done
  28.   -- Generating done
  29.   -- Build files have been written to: /cpp_inference_example/build
  30. $ make -j
  31.   Scanning dependencies of target custom_class
  32.   [ 25%] Building CXX object custom_class_project/CMakeFiles/custom_class.dir/class.cpp.o
  33.   [ 50%] Linking CXX shared library libcustom_class.so
  34.   [ 50%] Built target custom_class
  35.   Scanning dependencies of target infer
  36.   [ 75%] Building CXX object CMakeFiles/infer.dir/infer.cpp.o
  37.   [100%] Linking CXX executable infer
  38.   [100%] Built target infer

現(xiàn)在我們可以運行令人興奮的 C ++二進(jìn)制文件：

  1. $ ./infer
  2.   momfoobarbaz

難以置信！

定義自定義 C ++類的序列化/反序列化方法

如果您嘗試將具有自定義綁定 C ++類的ScriptModule保存為屬性，則會出現(xiàn)以下錯誤：

  1. # export_attr.py
  2. import torch
  5. torch.classes.load_library('libcustom_class.so')
  8. class Foo(torch.nn.Module):
  9.   def __init__(self):
  10.       super().__init__()
  11.       self.stack = torch.classes.Stack(["just", "testing"])
  14.   def forward(self, s : str) -> str:
  15.       return self.stack.pop() + s
  18. scripted_foo = torch.jit.script(Foo())
  21. scripted_foo.save('foo.pt')
  1. $ python export_attr.py
  2. 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)

這是因為 TorchScript 無法自動找出 C ++類中保存的信息。 您必須手動指定。 這樣做的方法是使用class_上的特殊def_pickle方法在類上定義__getstate__和__setstate__方法。

注意

TorchScript 中__getstate__和__setstate__的語義與 Python pickle 模塊的語義相同。 您可以閱讀更多有關(guān)如何使用這些方法的信息。

這是一個如何更新Stack類的注冊碼以包含序列化方法的示例：

  1. static auto testStack =
  2.   torch::jit::class_<Stack<std::string>>("Stack")
  3.       .def(torch::jit::init<std::vector<std::string>>())
  4.       .def("top", [](const c10::intrusive_ptr<Stack<std::string>>& self) {
  5.         return self->stack_.back();
  6.       })
  7.       .def("push", &Stack<std::string>::push)
  8.       .def("pop", &Stack<std::string>::pop)
  9.       .def("clone", &Stack<std::string>::clone)
  10.       .def("merge", &Stack<std::string>::merge)
  11.       // class_<>::def_pickle allows you to define the serialization
  12.       // and deserialization methods for your C++ class.
  13.       // Currently, we only support passing stateless lambda functions
  14.       // as arguments to def_pickle
  15.       .def_pickle(
  16.             // __getstate__
  17.             // This function defines what data structure should be produced
  18.             // when we serialize an instance of this class. The function
  19.             // must take a single `self` argument, which is an intrusive_ptr
  20.             // to the instance of the object. The function can return
  21.             // any type that is supported as a return value of the TorchScript
  22.             // custom operator API. In this instance, we've chosen to return
  23.             // a std::vector<std::string> as the salient data to preserve
  24.             // from the class.
  25.             [](const c10::intrusive_ptr<Stack<std::string>>& self)
  26.                 -> std::vector<std::string> {
  27.               return self->stack_;
  28.             },
  29.             // __setstate__
  30.             // This function defines how to create a new instance of the C++
  31.             // class when we are deserializing. The function must take a
  32.             // single argument of the same type as the return value of
  33.             // `__getstate__`. The function must return an intrusive_ptr
  34.             // to a new instance of the C++ class, initialized however
  35.             // you would like given the serialized state.
  36.             [](std::vector<std::string> state)
  37.                 -> c10::intrusive_ptr<Stack<std::string>> {
  38.               // A convenient way to instantiate an object and get an
  39.               // intrusive_ptr to it is via `make_intrusive`. We use
  40.               // that here to allocate an instance of Stack<std::string>
  41.               // and call the single-argument std::vector<std::string>
  42.               // constructor with the serialized state.
  43.               return c10::make_intrusive<Stack<std::string>>(std::move(state));
  44.             });

注意

我們采用與 pickle API 中的 pybind11 不同的方法。 pybind11 作為傳遞給class_::def()的特殊功能pybind11::pickle()，為此我們有一個單獨的方法def_pickle。 這是因為名稱torch::jit::pickle已經(jīng)被使用，我們不想引起混淆。

以這種方式定義(反）序列化行為后，腳本現(xiàn)在可以成功運行：

  1. import torch
  4. torch.classes.load_library('libcustom_class.so')
  7. class Foo(torch.nn.Module):
  8.     def __init__(self):
  9.         super().__init__()
  10.         self.stack = torch.classes.Stack(["just", "testing"])
  13.     def forward(self, s : str) -> str:
  14.         return self.stack.pop() + s
  17. scripted_foo = torch.jit.script(Foo())
  20. scripted_foo.save('foo.pt')
  21. loaded = torch.jit.load('foo.pt')
  24. print(loaded.stack.pop())
  1. $ python ../export_attr.py
  2. testing

結(jié)論

本教程向您介紹了如何向 TorchScript(以及擴展為 Python）公開 C ++類，如何注冊其方法，如何從 Python 和 TorchScript 使用該類以及如何使用該類保存和加載代碼以及運行該代碼。 在獨立的 C ++過程中。 現(xiàn)在，您可以使用與第三方 C ++庫接口的 C ++類擴展 TorchScript 模型，或?qū)崿F(xiàn)需要 Python，TorchScript 和 C ++之間的界線才能平滑融合的任何其他用例。

與往常一樣，如果您遇到任何問題或疑問，可以使用我們的論壇或 GitHub 問題進(jìn)行聯(lián)系。 另外，我們的常見問題解答(FAQ）頁面可能包含有用的信息。




以上內(nèi)容是否對您有幫助：
在線筆記
App下載
App下載

掃描二維碼

下載編程獅App

公眾號
微信公眾號

編程獅公眾號