使用自定义 C++ 类扩展 TorchScript

// 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<classT>
structMyStackClass:torch::CustomClassHolder{
std::vector<T>stack_;
MyStackClass(std::vector<T>init):stack_(init.begin(),init.end()){}

voidpush(Tx){
stack_.push_back(x);
}
Tpop(){
autoval=stack_.back();
stack_.pop_back();
returnval;
}

c10::intrusive_ptr<MyStackClass>clone()const{
returnc10::make_intrusive<MyStackClass>(stack_);
}

voidmerge(constc10::intrusive_ptr<MyStackClass>&c){
for(auto&elem:c->stack_){
push(elem);
}
}
};

// 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",[](constc10::intrusive_ptr<MyStackClass<std::string>>&self){
returnself->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_minimum_required(VERSION3.1FATAL_ERROR)
project(custom_class)

find_package(TorchREQUIRED)

# Define our library target
add_library(custom_classSHAREDclass.cpp)
set(CMAKE_CXX_STANDARD14)
# Link against LibTorch
target_link_libraries(custom_class"${TORCH_LIBRARIES}")

custom_class_project/
  class.cpp
  CMakeLists.txt
  build/

$cdbuild
$cmake-DCMAKE_PREFIX_PATH="$(python-c'import torch.utils; print(torch.utils.cmake_prefix_path)')"..
*-TheCcompileridentificationisGNU7.3.1
*-TheCXXcompileridentificationisGNU7.3.1
*-CheckforworkingCcompiler:/opt/rh/devtoolset-7/root/usr/bin/cc
*-CheckforworkingCcompiler:/opt/rh/devtoolset-7/root/usr/bin/cc--works
*-DetectingCcompilerABIinfo
*-DetectingCcompilerABIinfo-done
*-DetectingCcompilefeatures
*-DetectingCcompilefeatures-done
*-CheckforworkingCXXcompiler:/opt/rh/devtoolset-7/root/usr/bin/c++
*-CheckforworkingCXXcompiler:/opt/rh/devtoolset-7/root/usr/bin/c++--works
*-DetectingCXXcompilerABIinfo
*-DetectingCXXcompilerABIinfo-done
*-DetectingCXXcompilefeatures
*-DetectingCXXcompilefeatures-done
*-Lookingforpthread.h
*-Lookingforpthread.h-found
*-Lookingforpthread_create
*-Lookingforpthread_create-notfound
*-Lookingforpthread_createinpthreads
*-Lookingforpthread_createinpthreads-notfound
*-Lookingforpthread_createinpthread
*-Lookingforpthread_createinpthread-found
*-FoundThreads:TRUE
*-Foundtorch:/torchbind_tutorial/libtorch/lib/libtorch.so
*-Configuringdone
*-Generatingdone
*-Buildfileshavebeenwrittento:/torchbind_tutorial/build
$make-j
Scanningdependenciesoftargetcustom_class
[50%]BuildingCXXobjectCMakeFiles/custom_class.dir/class.cpp.o
[100%]LinkingCXXsharedlibrarylibcustom_class.so
[100%]Builttargetcustom_class

custom_class_project/
  class.cpp
  CMakeLists.txt
  build/
    libcustom_class.so

importtorch

# `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
defdo_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

importtorch

torch.classes.load_library('build/libcustom_class.so')


classFoo(torch.nn.Module):
    def__init__(self):
        super().__init__()

    defforward(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')

cpp_inference_example/
  infer.cpp
  CMakeLists.txt
  foo.pt
  build/
  custom_class_project/
    class.cpp
    CMakeLists.txt
    build/

#include<torch/script.h>

#include<iostream>
#include<memory>

intmain(intargc,constchar*argv[]){
torch::jit::Modulemodule;
try{
// Deserialize the ScriptModule from a file using torch::jit::load().
module=torch::jit::load("foo.pt");
}
catch(constc10::Error&e){
std::cerr<<"error loading the model\n";
return-1;
}

std::vector<c10::IValue>inputs={"foobarbaz"};
autooutput=module.forward(inputs).toString();
std::cout<<output->string()<<std::endl;
}

cmake_minimum_required(VERSION3.1FATAL_ERROR)
project(infer)

find_package(TorchREQUIRED)

add_subdirectory(custom_class_project)

# Define our library target
add_executable(inferinfer.cpp)
set(CMAKE_CXX_STANDARD14)
# 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-neededcustom_class)

$cdbuild
$cmake-DCMAKE_PREFIX_PATH="$(python-c'import torch.utils; print(torch.utils.cmake_prefix_path)')"..
*-TheCcompileridentificationisGNU7.3.1
*-TheCXXcompileridentificationisGNU7.3.1
*-CheckforworkingCcompiler:/opt/rh/devtoolset-7/root/usr/bin/cc
*-CheckforworkingCcompiler:/opt/rh/devtoolset-7/root/usr/bin/cc--works
*-DetectingCcompilerABIinfo
*-DetectingCcompilerABIinfo-done
*-DetectingCcompilefeatures
*-DetectingCcompilefeatures-done
*-CheckforworkingCXXcompiler:/opt/rh/devtoolset-7/root/usr/bin/c++
*-CheckforworkingCXXcompiler:/opt/rh/devtoolset-7/root/usr/bin/c++--works
*-DetectingCXXcompilerABIinfo
*-DetectingCXXcompilerABIinfo-done
*-DetectingCXXcompilefeatures
*-DetectingCXXcompilefeatures-done
*-Lookingforpthread.h
*-Lookingforpthread.h-found
*-Lookingforpthread_create
*-Lookingforpthread_create-notfound
*-Lookingforpthread_createinpthreads
*-Lookingforpthread_createinpthreads-notfound
*-Lookingforpthread_createinpthread
*-Lookingforpthread_createinpthread-found
*-FoundThreads:TRUE
*-Foundtorch:/local/miniconda3/lib/python3.7/site-packages/torch/lib/libtorch.so
*-Configuringdone
*-Generatingdone
*-Buildfileshavebeenwrittento:/cpp_inference_example/build
$make-j
Scanningdependenciesoftargetcustom_class
[25%]BuildingCXXobjectcustom_class_project/CMakeFiles/custom_class.dir/class.cpp.o
[50%]LinkingCXXsharedlibrarylibcustom_class.so
[50%]Builttargetcustom_class
Scanningdependenciesoftargetinfer
[75%]BuildingCXXobjectCMakeFiles/infer.dir/infer.cpp.o
[100%]LinkingCXXexecutableinfer
[100%]Builttargetinfer

$./infer
momfoobarbaz

# export_attr.py
importtorch

torch.classes.load_library('build/libcustom_class.so')


classFoo(torch.nn.Module):
    def__init__(self):
        super().__init__()
        self.stack = torch.classes.my_classes.MyStackClass(["just", "testing"])

    defforward(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())

$pythonexport_attr.py
RuntimeError:CannotserializecustomboundC++class__torch__.torch.classes.my_classes.MyStackClass.Pleasedefineserializationmethodsviadef_pickleforthisclass.(pushIValueImplat../torch/csrc/jit/pickler.cpp:128)

// 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.
[](constc10::intrusive_ptr<MyStackClass<std::string>>&self)
*>std::vector<std::string>{
returnself->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.
returnc10::make_intrusive<MyStackClass<std::string>>(std::move(state));
});

$python../export_attr.py
testing

c10::intrusive_ptr<MyStackClass<std::string>>manipulate_instance(constc10::intrusive_ptr<MyStackClass<std::string>>&instance){
instance->pop();
returninstance;
}

m.def(
"manipulate_instance(__torch__.torch.classes.my_classes.MyStackClass x) -> __torch__.torch.classes.my_classes.MyStackClass Y",
manipulate_instance
);

classTryCustomOp(torch.nn.Module):
    def__init__(self):
        super(TryCustomOp, self).__init__()
        self.f = torch.classes.my_classes.MyStackClass(["foo", "bar"])

    defforward(self):
        return torch.ops.my_classes.manipulate_instance(self.f)