C++加载torch-geometric导出的TorchScript模型报未知算子错误

• 需在C++环境加载Python环境训练生成的TorchScript模型，模型依赖pytorch-geometric、pytorch-scatter、pytorch-sparse、pytorch-cluster四个库，Python侧模型导出环节已验证正常。
• 已参考官方指引完成pytorch-scatter、pytorch-sparse、pytorch-cluster的C++ API安装，其中pytorch-scatter、pytorch-sparse链接运行正常，按相同方式配置pytorch-cluster链接后运行失败。
• 编译运行程序时报错如下：

Unknown builtin op: torch_cluster::fps.
Could not find any similar ops to torch_cluster::fps. This op may not exist or may not be currently supported in TorchScript.

CMakeLists.txt #

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(custom_ops)

find_package(TorchSparse REQUIRED)
find_package(TorchScatter REQUIRED)
find_package(TorchCluster REQUIRED)

add_executable(example-app example-app.cpp)

target_compile_features(example-app PUBLIC cxx_range_for)
target_link_libraries(example-app TorchSparse::TorchSparse)
target_link_libraries(example-app TorchScatter::TorchScatter)
target_link_libraries(example-app TorchCluster::TorchCluster)

set_property(TARGET example-app PROPERTY CXX_STANDARD 14)

example-app.cpp #

#include <torch/script.h> // One-stop header.
#include <torchsparse/sparse.h>
#include <torchscatter/scatter.h>
#include <torchcluster/cluster.h>
#include <iostream>
#include <memory>

int main(int argc, const char* argv[]) {
  if (argc != 2) {
    std::cerr << "usage: example-app <path-to-exported-script-module>\n";
    return -1;
  }


  torch::jit::script::Module module;
  try {
    // 从文件反序列化ScriptModule
    module = torch::jit::load(argv[1]);
  }
  catch (const c10::Error& e) {
      std::cerr << e.what();
    std::cerr << "error loading the model\n";
    return -1;
  }

  std::cout << "ok\n";
}

按以下优先级依次排查：

1. 补全fps算子的头文件引入
   多数版本的pytorch-cluster C++库不会在引入torchcluster/cluster.h时自动注册所有TorchScript算子，fps属于独立实现的算子模块，需要在cpp文件的include区域补充对应头文件：

   #include <torchcluster/fps.h>
   

   修改后重新编译，确认链接阶段将fps算子的符号打包进可执行文件。
2. 对齐Python侧与C++侧的依赖版本
   版本不匹配是自定义算子加载失败的最高频原因：
   • 在Python环境执行命令记录PyTorch、torch_cluster的版本号与编译CUDA版本：

     import torch, torch_cluster
     print(torch.__version__, torch.version.cuda, torch_cluster.__version__)
     

   • 核对C++侧find_package找到的Torch、TorchCluster版本，二者必须完全一致，小版本差异、CUDA版本差异都会导致算子schema不匹配、注册失败。
3. 调整CMake链接配置
   部分编译器下静态库链接顺序会影响符号加载，同时动态库依赖需要配置rpath避免运行时找不到依赖，修改CMakeLists.txt的链接配置如下：

   # 调整链接顺序，被依赖的库放在更靠后的位置
   target_link_libraries(example-app 
     TorchCluster::TorchCluster
     TorchScatter::TorchScatter
     TorchSparse::TorchSparse
     "${TORCH_LIBRARIES}"
   )
   # 配置rpath，避免运行时找不到各依赖的动态库
   set_target_properties(example-app PROPERTIES
     BUILD_RPATH "${TorchCluster_DIR}/../lib;${TorchScatter_DIR}/../lib;${TorchSparse_DIR}/../lib"
     CXX_STANDARD 14
   )
   

4. 校验算子注册状态
   编译完成后，在torch::jit::load调用前添加如下代码，打印所有已注册的torch_cluster域算子：

   for (const auto& op : torch::jit::getAllOperators()) {
     if (op->schema().name().rfind("torch_cluster::", 0) == 0) {
       std::cout << "Registered PyG cluster op: " << op->schema().name() << std::endl;
     }
   }
   

   • 如果运行输出中没有torch_cluster::fps，说明链接阶段没有将fps算子实现链接进可执行文件，回到前两步检查头文件引入、版本匹配、链接配置。
   • 如果输出中存在torch_cluster::fps但加载模型仍报错，说明模型导出时的算子schema和当前C侧算子schema不匹配，在和C依赖版本完全一致的Python环境中重新导出TorchScript模型即可。

内容的提问来源于stack exchange，提问作者Yoann PITARCH