Table of Contents
  • Tutorials >
  • 커스텀 C++ 클래스로 TorchScript 확장하기
Shortcuts
advanced/torch_script_custom_classes

커스텀 C++ 클래스로 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::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++ code를 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를 호출한 다음 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: /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

이제 무엇보다도 빌드 디렉토리에 동적 라이브러리 파일이 있다는 것을 알게 될 것입니다. 리눅스에서는 아마도 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 코드 저장, 읽기 및 실행

libtorch를 사용하여 C++ 프로세스에서 커스텀 등록 C++ 클래스를 사용할 수도 있습니다. 예를 들어 MyStackClass 클래스에서 메소드를 인스턴스화하고 호출하는 간단한 nn.Module 을 정의해 보겠습니다:

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 프로그램을 포함합니다.

이제 이 모델과 필요한 .so 파일을 읽어들이는 방법을 보여주기 위해 새 CMake 프로젝트를 정의하겠습니다. 이 작업을 수행하는 방법에 대한 자세한 내용은 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로/에서 이동

TorchScript 메소드에서 IValue 를 가져오거나 반환하기, 또는 C++에서 커스텀 클래스 속성을 인스턴스화하려는 경우와 같이 커스텀 클래스를 IValue 안팎으로 이동해야 할 수도 있습니다. 커스텀 C++ 클래스 인스턴스에서 IValue 를 생성하려면:

  • torch::make_custom_class<T>() 는 제공하는 인수 집합을 사용하고 해당 인수 집합과 일치하는 T의 생성자를 호출하며 해당 인스턴스를 래핑하고 반환하는 c10::intrusive_ptr<T>와 유사한 API를 제공합니다. 그러나 커스텀 클래스 객체에 대한 포인터만 반환하는 대신 객체를 래핑하는 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++ 클래스에서 저장한 정보를 자동으로 파악할 수 없기 때문입니다. 수동으로 지정해야 합니다. 그렇게 하는 방법은 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이 class_::def() 로 전달되는 특수 함수 pybind11::pickle() 인 반면, 이 목적을 위해 별도의 메소드 def_pickle 를 가지고 있습니다. 이미 torch::jit::pickle 라는 이름이 사용되었고 혼동을 일으키고 싶지 않았기 때문입니다.

이러한 방식으로 (역)직렬화 동작을 정의하면 이제 스크립트를 성공적으로 실행할 수 있습니다:

$ python ../export_attr.py
testing

바인딩된 C++ 클래스를 사용하거나 반환하는 커스텀 연산자 정의

커스텀 C++ 클래스를 정의한 후에는 해당 클래스를 인수로 사용하거나 커스텀 연산자(예를 들어 free 함수)에서 반환할 수도 있습니다. 다음과 같은 free 함수가 있다고 가정합니다:

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에 대한 자세한 내용은 커스텀 C++ 연산자로 TorchScript 확장 을 참조하세요.

이 작업이 완료되면 다음 예제와 같이 연산자를 사용할 수 있습니다:

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++ 클래스를 인수로 사용하는 연산자를 등록하려면 커스텀 클래스가 이미 등록되어 있어야 합니다. 커스텀 클래스 등록과 free 함수 정의가 동일한 TORCH_LIBRARY 블록에 있고 커스텀 클래스 등록이 먼저 오게 하여 이를 시행할 수 있습니다. 향후 어떤 순서로든 등록할 수 있도록 이 요구사항을 완화할 수 있습니다.

결론

이 튜토리얼에서는 독립된 C++ 프로세스에서 C++ 클래스를 TorchScript 및 확장 Python에 나타내는 방법, 해당 메소드를 등록하는 방법, Python 및 TorchScript에서 해당 클래스를 사용하는 방법, 클래스를 사용하여 코드를 저장 및 읽어들이고 해당 코드를 실행하는 방법을 안내했습니다. 이제 타사 C++ 라이브러리와 인터페이스가 있는 C++ 클래스로 TorchScript 모델을 확장하거나, Python, TorchScript 및 C++ 간의 라인이 원활하게 혼합되어야 하는 다른 사용 사례를 구현할 준비가 되었습니다.

언제나처럼 문제를 마주치거나 질문이 있으면 저희 forum 또는 GitHub issues 에 올려주시면 되겠습니다. 또한 자주 묻는 질문(FAQ) 페이지 에 유용한 정보가 있을 수 있습니다.

더 궁금하시거나 개선할 내용이 있으신가요? 커뮤니티에 참여해보세요!
한국어 커뮤니티 바로가기
이 튜토리얼이 어떠셨나요? 평가해주시면 이후 개선에 참고하겠습니다! :)

© Copyright 2018-2024, PyTorch & 파이토치 한국 사용자 모임(PyTorch Korea User Group).

Built with Sphinx using a theme provided by Read the Docs.

PyTorchKorea @ GitHub

파이토치 한국 사용자 모임을 GitHub에서 만나보세요.

GitHub로 이동

한국어 튜토리얼

한국어로 번역 중인 PyTorch 튜토리얼입니다.

튜토리얼로 이동

커뮤니티

다른 사용자들과 의견을 나누고, 도와주세요!

커뮤니티로 이동