Question

我想学习修改TensorFlow内置运算符内核的最佳方法。例如，我想在static const double A中修改tensorflow/core/kernels/resize_bicubic_op.cc的值。我想出了两种可能的方法：

直接修改它并重新编译整个TensorFlow库。该解决方案的问题是：A。这影响使用双三次插值的所有函数。 B.这需要我重新编译整个库，并且在从二进制文件安装时不起作用。
将其定义为自定义操作。问题是在源代码中，内部没有REGISTER_OP()。我不知道如何为这个双三次函数编写REGISTER_OP()以及是否需要进行其他修改。

还有其他更好的方法吗？

感谢。

Answer 1

解决此问题的最佳方法是构建自定义操作。有关如何添加自定义操作的详细信息，请参阅this tutorial。对tf.image.resize_bicubic()操作的REGISTER_OP调用位于tensorflow/core/ops/image_ops.cc。

另一种方法是重新使用相同的 op 注册，并使用替代实现注册新的内核。这样您就可以使用（实验性）Graph.kernel_label_map() API为"ResizeBicubic"操作选择替代实现。例如，您可以在Python程序中执行以下操作：

_ = tf.load_op_library(...)  # Load the .so containing your implementation.

with tf.get_default_graph().kernel_label_map({"ResizeBicubic": "my_impl"}):
  images = tf.image.resize_bicubic(...)  # Will use your implementation.

...并添加一个内核注册，用您的C ++代码指定标签"my_impl"：

template <typename Device, typename T>
class MyResizeBicubicOp<Device, T> : public OpKernel {
  // Custom implementation goes here...
}

#define REGISTER_KERNEL(T)                            \
  REGISTER_KERNEL_BUILDER(Name("ResizeBicubic")       \
                              .Device(DEVICE_CPU)     \
                              .Label("my_impl")       \
                              .TypeConstraint<T>("T") \
                              .HostMemory("size"),    \
                          MyResizeBicubicOp<CPUDevice, T>);

TF_CALL_REAL_NUMBER_TYPES(REGISTER_KERNEL);

修改内置TensorFlow内核的最佳方法

1 个答案: