• Tensorflow-serving支持pytorch模型的推理
  • 背景
  • 步骤
    • 1. 新增tfserving third_party 的依赖
    • 2. 新增Servable
    • 3. 注册Pytorch的Servable
    • 4. 复用TF-serving的gRPC协议
  • 踩坑记
  • Github 地址

    Tensorflow-serving支持pytorch模型的推理

    背景

    将 Pytorch 的模型使用 TF-Serving 框架去做推理。有几个理由：

1. Pytorch 至今没有官方的工业级推理框架(截止本文前)
2. TF-Serving 的可扩展性不错，根据官方文档是可以自定义推理模块的
3. 当有多个业务方采用不同框架训练模型，同时存在依赖调用时，可以将不同模型的结果在内部TF-Serving内部做级联调用，降低多个推理框架之间调用的成本。

思路：

1. 直接编译 Pytorch 的推理库，然后集成到 TF-Serving 中。
2. 扩展TF-Serving的Servable，将Pytorch的推理库封装成一个Servable。

   步骤

   1. 新增tfserving third_party 的依赖

   在 tensorflow_serving/workspace.bzl： ```

• ===== libpytorch dependencies =====

• http_archive(
• name = “com_local_libtorch”,
• urls = [
• “http://127.0.0.1:8000/libtorch-cxx11-abi-shared-with-deps-1.4.0.zip“,
• ],
• sha256 = “f214bfde532877aa5d4e0803e51a28fa8edd97b6a44b6615f75a70352b6b542e”,
• strip_prefix = “libtorch-cxx11-abi-shared-with-deps-1.4.0”,
• build_file = “@//third_party/libtorch:BUILD”,
• )

  同时在 third_party 里指定要编译的BUILD

  root@4bf9f1d29891:/tensorflow-serving# cat third_party/libtorch/BUILD

  libtorch (libtorch.org) library.

  package( default_visibility = [“//visibility:public”], ) licenses([“notice”]) # BSD exports_files([“LICENSE”]) include_files = [ “libtorch/include/*/.h”, ] cc_import( name = “libtorch”, hdrs = include_files, static_library = “libtorch/lib/.a”, shared_library = “libtorch/lib/.so”, )

  <a name="S80Nh"></a>
  #### 2. 新增Servable
  Servable 需要对应的 Loader 和 SourceAdapter。而 SourceAdapter 又是通过 SourceAdapterConfig 创建的，SourceAdapterConfig 是作为插件插在 PlatformConfigMap 里的，所以这里的先创建一个 SourceAdapter.proto，里面定义SourceAdapterConfig。<br />相关的配置文件:<br />options.model_server_config

  { “model_base_name” : “default”, “model_platform” : “Tensorflow”, } options.platform_config_map default is : TensorFlowPlatformConfigMap // Configuration for a servable platform e.g. tensorflow or other ML systems. message PlatformConfig { // The config proto for a SourceAdapter in the StoragePathSourceAdapter // registry. google.protobuf.Any source_adapter_config = 1; }; message PlatformConfigMap { // A map from a platform name to a platform config. The platform name is used // in ModelConfig.model_platform. map platform_configs = 1; }; ReloadConfig(model_server_config) { }

  <a name="WGuyL"></a>
  #### 3. 注册Pytorch的Servable

  // Registers a factory that creates subclasses of BaseClass by calling // ClassCreator::Create().

  define REGISTER_CLASS(RegistryName, BaseClass, ClassCreator, config_proto, \

                ...)                                                 \

  REGISTERCLASSUNIQ_HELPER(__COUNTER, RegistryName, BaseClass, \

                      ClassCreator, config_proto, ##__VA_ARGS__)
  

  define REGISTER_STORAGE_PATH_SOURCE_ADAPTER(ClassCreator, ConfigProto) \

  REGISTER_CLASS(StoragePathSourceAdapterRegistry, StoragePathSourceAdapter, \

          ClassCreator, ConfigProto);
  

  在调用的时候如下:<br />StoragePathSourceAdapterRegistry:注册器名称;<br />StoragePathSourceAdapter:父类;<br />SavedModelBundleSourceAdapterCreator:子类创建器;<br />SavedModelBundleSourceAdapterConfig:对应配置文件;
  

  REGISTER_STORAGE_PATH_SOURCE_ADAPTER(SavedModelBundleSourceAdapterCreator, SavedModelBundleSourceAdapterConfig); ```

  4. 复用TF-serving的gRPC协议

  需要做2件事：

1. pytorch客户端将预测数据使用protobuf序列化，在server端需要反序列化之后，转化为pytorch的Tensor格式。
   
2. 同理，pytorch的lib库推理结束之后，也需要将结果序列化为protobuf的格式回传。
   

因此涉及2个功能点：TorchTensor2Proto、Proto2TorchTensor

踩坑记

1. static 变量不初始化的问题：
   cc_library() 要标记为 alwayslink=1
   
2. 日志不输出的问题：
   如果同时需要包含 torch/xxx.h，需保证 torch/xxx.h 的include 语句在所有的 #include tensorflow 之前
   
3. 单独的 test cc_binaray() 必须要连接多线程的库
   -lpthread
   
4. -D_GLIBCXX_USE_CXX11_ABI 不兼容的问题：
   如果 libtorch 使用 -D_GLIBCXX_USE_CXX11_ABI=1 编译的，则必须保证 tensorflow_model_server 的编译也是用相同的选项，而且是全局的，需要在 .bazelrc 里加入
   build --cxxopt=-D_GLIBCXX_USE_CXX11_ABI=1。
   
5. 打开GDB 调试选项的问题：
   编译时不指定-g 是无法使用gdb 调试的
   build --cxxopt=-g
   

   Github 地址