服务于自定义 gpu op 的 Tensorflow 在编译时找不到依赖项

Question

我参考了有关为 tensorflow 制作自定义 gpu op 的指南，并且可以制作共享库。对于 tensorflow-serving，我调整了所需的路径，但在构建时出现错误：

   ERROR: /home/g360/Documents/eduardss/serving/tensorflow_serving/custom_ops/CUSTOM_OP/BUILD:32:1: undeclared inclusion(s) in rule '//tensorflow_serving/custom_ops/CUSTOM_OP:CUSTOM_OP_ops_gpu':
this rule is missing dependency declarations for the following files included by 'tensorflow_serving/custom_ops/CUSTOM_OP/cc/magic_op.cu.cc':
  'external/org_tensorflow/tensorflow/core/platform/stream_executor.h'
  'external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_platform_id.h'
  'external/org_tensorflow/tensorflow/stream_executor/platform.h'
  'external/org_tensorflow/tensorflow/stream_executor/device_description.h'
  'external/org_tensorflow/tensorflow/stream_executor/launch_dim.h'
  'external/org_tensorflow/tensorflow/stream_executor/platform/port.h'
  'external/org_tensorflow/tensorflow/stream_executor/device_options.h'
  'external/org_tensorflow/tensorflow/stream_executor/platform/logging.h'
  'external/org_tensorflow/tensorflow/stream_executor/lib/status.h'
  'external/org_tensorflow/tensorflow/stream_executor/lib/error.h'
  'external/org_tensorflow/tensorflow/stream_executor/lib/status_macros.h'
  'external/org_tensorflow/tensorflow/stream_executor/lib/statusor.h'
  'external/org_tensorflow/tensorflow/stream_executor/lib/statusor_internals.h'
  'external/org_tensorflow/tensorflow/stream_executor/plugin.h'
  'external/org_tensorflow/tensorflow/stream_executor/trace_listener.h'
  'external/org_tensorflow/tensorflow/stream_executor/device_memory.h'
  'external/org_tensorflow/tensorflow/stream_executor/kernel.h'
  'external/org_tensorflow/tensorflow/stream_executor/kernel_cache_config.h'
  'external/org_tensorflow/tensorflow/stream_executor/lib/array_slice.h'
  'external/org_tensorflow/tensorflow/stream_executor/dnn.h'
  'external/org_tensorflow/tensorflow/stream_executor/event.h'
  'external/org_tensorflow/tensorflow/stream_executor/host/host_platform_id.h'
  'external/org_tensorflow/tensorflow/stream_executor/multi_platform_manager.h'
  'external/org_tensorflow/tensorflow/stream_executor/lib/initialize.h'
  'external/org_tensorflow/tensorflow/stream_executor/platform/initialize.h'
  'external/org_tensorflow/tensorflow/stream_executor/platform/platform.h'
  'external/org_tensorflow/tensorflow/stream_executor/platform/default/initialize.h'
  'external/org_tensorflow/tensorflow/stream_executor/platform/dso_loader.h'
  'external/org_tensorflow/tensorflow/stream_executor/platform/default/dso_loader.h'
  'external/org_tensorflow/tensorflow/stream_executor/rocm/rocm_platform_id.h'
  'external/org_tensorflow/tensorflow/stream_executor/scratch_allocator.h'
  'external/org_tensorflow/tensorflow/stream_executor/temporary_device_memory.h'
  'external/org_tensorflow/tensorflow/stream_executor/stream.h'
  'external/org_tensorflow/tensorflow/stream_executor/blas.h'
  'external/org_tensorflow/tensorflow/stream_executor/host_or_device_scalar.h'
  'external/org_tensorflow/tensorflow/stream_executor/fft.h'
  'external/org_tensorflow/tensorflow/stream_executor/platform/thread_annotations.h'
  'external/org_tensorflow/tensorflow/stream_executor/temporary_memory_manager.h'
  'external/org_tensorflow/tensorflow/stream_executor/stream_executor.h'
  'external/org_tensorflow/tensorflow/stream_executor/kernel_spec.h'
  'external/org_tensorflow/tensorflow/stream_executor/stream_executor_pimpl.h'
  'external/org_tensorflow/tensorflow/stream_executor/device_memory_allocator.h'
  'external/org_tensorflow/tensorflow/stream_executor/lib/threadpool.h'
  'external/org_tensorflow/tensorflow/stream_executor/lib/env.h'
  'external/org_tensorflow/tensorflow/stream_executor/lib/thread_options.h'
  'external/org_tensorflow/tensorflow/stream_executor/rng.h'
  'external/org_tensorflow/tensorflow/stream_executor/shared_memory_config.h'
  'external/org_tensorflow/tensorflow/stream_executor/stream_executor_internal.h'
  'external/org_tensorflow/tensorflow/stream_executor/allocator_stats.h'
  'external/org_tensorflow/tensorflow/stream_executor/module_spec.h'
  'external/org_tensorflow/tensorflow/stream_executor/plugin_registry.h'
  'external/org_tensorflow/tensorflow/stream_executor/timer.h'

有问题的函数导入：

#if GOOGLE_CUDA

#define EIGEN_USE_GPU

#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
#include "tensorflow/core/util/gpu_kernel_helper.h"

#include "math.h"

#include <iostream>

在这些依赖项中，我认为 gpu_kernel_helper.h 是导致错误的原因，因为我的 BUILD 文件依赖项是：

    deps = [
    "@org_tensorflow//tensorflow/core:framework",
    "@org_tensorflow//tensorflow/core:lib",
    "@org_tensorflow//third_party/eigen3",
] + if_cuda_is_configured([":cuda",  "@local_config_cuda//cuda:cuda_headers"]),

如果我尝试直接导入它们，它会抱怨路径上没有 BUILD 文件。我不太熟悉 bazel 的构建过程，所以不明白它究竟需要如何 link 导入。

编辑 1 我使用了 tensorflow-serving 2.1.0 和 tensorflow/serving:2.1.0-devel-gpu docker 图像。查看@org_tensorflow/tensorflow/core/BUILD 实际上有一些参考 'gpu_kernel_helper.h':

tf_cuda_library(
name = "framework",
hdrs = [
       ...
       "util/gpu_kernel_helper.h",
       ]

但显然下游 link 缺少一些？

Answer 1

解决方法： 缺失的依赖可以通过以下方式链接：

"@org_tensorflow//tensorflow/core:stream_executor_headers_lib"

服务于自定义 gpu op 的 Tensorflow 在编译时找不到依赖项

Tensorflow serving custom gpu op cannot find dependency when compiling

tensorflow

bazel

tensorflow-serving