PyTorch 2.2 中文官方教程(十二)(3)

简介: 流畅的 Python 第二版(GPT 重译)(十)

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

译者:飞龙

协议:CC BY-NC-SA 4.0

本教程是自定义运算符教程的后续,介绍了我们为将 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 buildcmake,然后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++之间平滑过渡的用例。

相关文章
|
11月前
|
存储 物联网 PyTorch
基于PyTorch的大语言模型微调指南:Torchtune完整教程与代码示例
**Torchtune**是由PyTorch团队开发的一个专门用于LLM微调的库。它旨在简化LLM的微调流程,提供了一系列高级API和预置的最佳实践
549 59
基于PyTorch的大语言模型微调指南:Torchtune完整教程与代码示例
|
11月前
|
并行计算 监控 搜索推荐
使用 PyTorch-BigGraph 构建和部署大规模图嵌入的完整教程
当处理大规模图数据时,复杂性难以避免。PyTorch-BigGraph (PBG) 是一款专为此设计的工具,能够高效处理数十亿节点和边的图数据。PBG通过多GPU或节点无缝扩展,利用高效的分区技术,生成准确的嵌入表示,适用于社交网络、推荐系统和知识图谱等领域。本文详细介绍PBG的设置、训练和优化方法,涵盖环境配置、数据准备、模型训练、性能优化和实际应用案例,帮助读者高效处理大规模图数据。
205 5
|
并行计算 Ubuntu PyTorch
Ubuntu下CUDA、Conda、Pytorch联合教程
本文是一份Ubuntu系统下安装和配置CUDA、Conda和Pytorch的教程,涵盖了查看显卡驱动、下载安装CUDA、添加环境变量、卸载CUDA、Anaconda的下载安装、环境管理以及Pytorch的安装和验证等步骤。
2698 1
Ubuntu下CUDA、Conda、Pytorch联合教程
|
PyTorch 算法框架/工具 异构计算
PyTorch 2.2 中文官方教程(十九)(1)
PyTorch 2.2 中文官方教程(十九)
238 1
PyTorch 2.2 中文官方教程(十九)(1)
|
机器学习/深度学习 PyTorch 算法框架/工具
PyTorch 2.2 中文官方教程(十八)(4)
PyTorch 2.2 中文官方教程(十八)
210 1
|
PyTorch 算法框架/工具 异构计算
PyTorch 2.2 中文官方教程(二十)(4)
PyTorch 2.2 中文官方教程(二十)
295 0
PyTorch 2.2 中文官方教程(二十)(4)
|
Android开发 PyTorch 算法框架/工具
PyTorch 2.2 中文官方教程(二十)(2)
PyTorch 2.2 中文官方教程(二十)
264 0
PyTorch 2.2 中文官方教程(二十)(2)
|
iOS开发 PyTorch 算法框架/工具
PyTorch 2.2 中文官方教程(二十)(1)
PyTorch 2.2 中文官方教程(二十)
225 0
PyTorch 2.2 中文官方教程(二十)(1)
|
PyTorch 算法框架/工具 异构计算
PyTorch 2.2 中文官方教程(十九)(3)
PyTorch 2.2 中文官方教程(十九)
246 0
PyTorch 2.2 中文官方教程(十九)(3)
|
异构计算 PyTorch 算法框架/工具
PyTorch 2.2 中文官方教程(十九)(2)
PyTorch 2.2 中文官方教程(十九)
222 0
PyTorch 2.2 中文官方教程(十九)(2)

热门文章

最新文章

推荐镜像

更多