Help with custom c++ class in a jit model

C++
1

Hi all,
I am having some difficulties getting my custom c++ class to work correctly with torch::jit::load despite being able to use it in python.

Essentially I have a custom class that follows this structure in the cpp file

TORCH_LIBRARY(mycustomop, m) {
  m.class_<CustomComputer>("CustomComputer")
      .def(torch::init< Tensor, Tensor, Tensor, Tensor >())
      .def_pickle(
          // __getstate__
          [](const c10::intrusive_ptr<CustomComputer>& self) -> std::vector<Tensor> {
            std::vector<Tensor> state;
            state.push_back(self-> params.a);
            state.push_back(self->params.b);
            state.push_back(self->params.c);
            state.push_back(self->params.d);
            return state;
          },
          // __setstate__
          [](std::vector<Tensor> state) -> c10::intrusive_ptr<CustomComputer> {
            return c10::make_intrusive<CustomComputer>(
                state[0],
                state[1],
                state[2],
                state[3]);
          });

however when I attempt to jit load a model that uses this class I run into the follow exception

Exception:
Unknown type name '__torch__.torch.classes.mycustomop.CustomComputer':

I have no problem using this model on the python side if I do something like

import torch
import mycustomop
model = torch.jit.load('my_compiled_model')

Can someone help with how to properly register this custom class on the c++ side.

2

Hello @dma_science

I was able to reproduce your error on my system:
Unknown type name '__torch__.torch.classes.mycustomop.CustomComputer':

I managed to fix it and sucessfully load the model by linking the CustomComputer class during the build of the C++ binaries. In my case, when using CMake, adding the .cpp file with the CustomComputer implementation to the executables allows me to correctly load the model.

Maybe you can share exactly how you’re building your binaries, so we can take a closer look at where’s the issue.

Machine Learning Engineer at RidgeRun.ai
Contact us: support@ridgerun.ai

3

Thank you for the response! My custom class is being used as part of a larger project so the cmakelists.txt is rather long and complicated, i.e., it likely is a linking problem. Is it possible for you to show the solution for your minimal example? That will probably allow me to debug on my end. Thank you!

4

Sure, no problem. This is my CMakeLists.txt, where computer.cpp contains the implementation of the CustomComputer class as well as the TORCH_LIBRARY definition, which is the same as the one you shared earlier. main.cpp the is where model is loaded with torch::jit::load.

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(MyProject)

set(CMAKE_PREFIX_PATH "${CMAKE_CURRENT_SOURCE_DIR}/libtorch")
find_package(Torch REQUIRED)

add_executable(main main.cpp computer.cpp)
target_link_libraries(main "${TORCH_LIBRARIES}")
set_property(TARGET main PROPERTY CXX_STANDARD 17)

I managed to reproduce the error by removing computer.cpp from the add_executable line, which led me to believe that it stems from linking issues. Hope this helps.