10

为了学习如何编写自定义 TensorFlow 操作，我遵循了添加新操作教程并制作了一个“add_b”操作，b为每个输入值添加了一个标量。

add_b_op.cc：

#define EIGEN_USE_THREADS

#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"

#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/shape_inference.h"

using namespace tensorflow;

REGISTER_OP("AddB")
    .Attr("T: {float, double}")
    .Input("input: T")
    .Input("b: T")
    .Output("output: T")
    .SetShapeFn([] (shape_inference::InferenceContext* c) -> Status {
      shape_inference::ShapeHandle out;
      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &out));
      return shape_inference::UnchangedShape(c);
    })
//----------------------------------------------------------------------
    .Doc(R"doc(
Adds `b` to each input.

input: The input values.
b: A number to add to each input value.
)doc");


template <typename T>
class AddBCpuOp : public OpKernel {
 public:
  explicit AddBCpuOp(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    const Tensor& input_tensor = context->input(0);
    const auto input = input_tensor.flat<T>();

    Tensor* output_tensor = nullptr;
    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
                                                     &output_tensor));
    auto output = output_tensor->flat<T>();

    const Eigen::ThreadPoolDevice& d = context->eigen_device<Eigen::ThreadPoolDevice>();

    // Note: The mistake of adding 1 instead of `b` is intentional to be able to distinguish
    // the CPU and GPU implementations.
    output.device(d) = input + static_cast<T>(1);
  }
};

REGISTER_KERNEL_BUILDER(
    Name("AddB")
    .Device(DEVICE_CPU)
    .TypeConstraint<float>("T"),
    AddBCpuOp<float>);
REGISTER_KERNEL_BUILDER(
    Name("AddB")
    .Device(DEVICE_CPU)
    .TypeConstraint<double>("T"),
    AddBCpuOp<double>);


#if GOOGLE_CUDA

template <typename T>
bool LaunchAddBKernel(const T *__restrict__ d_input, int n, const T *__restrict__ d_b, T *__restrict__ d_output);

template <typename T>
class AddBGpuOp : public OpKernel {
 public:
  explicit AddBGpuOp(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    const Tensor& input_tensor = context->input(0);
    const auto input = input_tensor.flat<T>();

    const Tensor& b_tensor = context->input(1);
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(b_tensor.shape()),
                errors::InvalidArgument("add_b expects a scalar for `b`."));
    const auto b = b_tensor.scalar<T>();

    Tensor* output_tensor = nullptr;
    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
                                                     &output_tensor));
    auto output = output_tensor->flat<T>();

    OP_REQUIRES(context, LaunchAddBKernel(input.data(), input.dimension(0), b.data(), output.data()),
                errors::Internal("add_b: LaunchAddBKernel() failed."));
  }
};

REGISTER_KERNEL_BUILDER(
    Name("AddB")
    .Device(DEVICE_GPU)
    .TypeConstraint<float>("T"),
    AddBGpuOp<float>);
REGISTER_KERNEL_BUILDER(
    Name("AddB")
    .Device(DEVICE_GPU)
    .TypeConstraint<double>("T"),
    AddBGpuOp<double>);

#endif // if GOOGLE_CUDA

add_b_op.cu.cc

template <typename T, int BLOCK_DIM_X>
__global__ void AddBKernel(const T *__restrict__ d_input, int n, const T *__restrict__ d_b, T *__restrict__ d_output) {
  const int i = blockIdx.x * BLOCK_DIM_X + threadIdx.x;
  if (i < n) {
    d_output[i] = d_input[i] + *d_b;
  }
}

template <typename T>
bool LaunchAddBKernel(const T *__restrict__ d_input, int n, const T *__restrict__ d_b, T *__restrict__ d_output) {
  if (n <= 0) return true;

  constexpr int BLOCK_DIM_X = 256;
  AddBKernel<T, BLOCK_DIM_X><<<n / BLOCK_DIM_X + (n % BLOCK_DIM_X != 0), BLOCK_DIM_X>>>(d_input, n, d_b, d_output);
  return true;
}

// Explicit instantiations.
template bool LaunchAddBKernel<float>(const float *__restrict__, int, const float *__restrict__, float *__restrict__);
template bool LaunchAddBKernel<double>(const double *__restrict__, int, const double *__restrict__, double *__restrict__);

我故意在 CPU 实现中引入了一个错误，以便能够区分正在使用的是 CPU 还是 GPU 实现。

当我测试我的自定义操作时：

from __future__ import print_function
import tensorflow as tf

module = tf.load_op_library('custom_ops.so')
with tf.Session(config = tf.ConfigProto(log_device_placement = True)):
  print(module.add_b([5., 4., 3., 2., 1.], 8.).eval())

我得到以下输出：

我 tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:892] OS X 不支持 NUMA - 返回 NUMA 节点零
我 tensorflow/core/common_runtime/gpu/gpu_device.cc:951] 找到具有属性的设备 0：
名称：GeForce GT 750M
主要：3 次要：0 memoryClockRate (GHz) 0.9255
pciBusID 0000:01:00.0
总内存：2.00GiB
可用内存：1.80GiB
我 tensorflow/core/common_runtime/gpu/gpu_device.cc:972] DMA: 0
我 tensorflow/core/common_runtime/gpu/gpu_device.cc:982] 0: Y
I tensorflow/core/common_runtime/gpu/gpu_device.cc:1041] 创建 TensorFlow 设备 (/gpu:0) -> （设备：0，名称：GeForce GT 750M，pci 总线 ID：0000:01:00.0）
设备映射：
/job:localhost/replica:0/task:0/gpu:0 -> 设备：0，名称：GeForce GT 750M，pci 总线 ID：0000:01:00.0
我 tensorflow/core/common_runtime/direct_session.cc:252] 设备映射：
/job:localhost/replica:0/task:0/gpu:0 -> 设备：0，名称：GeForce GT 750M，pci 总线 ID：0000:01:00.0

AddB: /job:localhost/replica:0/task:0/gpu:0
我 tensorflow/core/common_runtime/simple_placer.cc:819] AddB: /job:localhost/replica:0/task:0/gpu:0
AddB/b: /job:localhost/replica:0/task:0/gpu:0
我 tensorflow/core/common_runtime/simple_placer.cc:819] AddB/b: /job:localhost/replica:0/task:0/gpu:0
AddB/输入：/job:localhost/replica:0/task:0/gpu:0
我 tensorflow/core/common_runtime/simple_placer.cc:819] AddB/input: /job:localhost/replica:0/task:0/gpu:0
[6. 5. 4. 3. 2.]

“设备放置日志”似乎表明该操作正在 GPU 上执行，但输出表明正在使用 CPU 实现。

当我注释掉用于DEVICE_CPU实现、重新编译和重新测试的两个 REGISTER_KERNEL_BUILDER() 注册时，我得到了预期的输出[ 13. 12. 11. 10. 9.]，但是有一个错误：

E tensorflow/core/common_runtime/executor.cc:334] 执行器创建内核失败。未找到：没有为与节点 AddB 兼容的 CPU 设备注册的“AddB”OpKernel = AddB[T=DT_FLOAT, _device="/job:localhost/replica:0/task:0/gpu:0"](AddB/input, AddB /b)
    . 注册：设备='GPU'；[DT_FLOAT] 中的 T
  设备='GPU'；[DT_DOUBLE] 中的 T

     [[节点：AddB = AddB[T=DT_FLOAT, _device="/job:localhost/replica:0/task:0/gpu:0"](AddB/input, AddB/b)]]

该错误消息对我来说似乎是一个错误，因为尽管错误显示“执行程序无法创建内核”，但显然创建了一个内核来在 GPU 上运行操作。

为什么使用 CPU 实现而不是 GPU 实现？

如果这很重要，以下是有关我的开发设置的详细信息：

  • 我正在使用带有内置 NVIDIA GeForce GT 750M（CUDA Compute Capability 3.0）的 MacBook Pro。
  • macOS Sierra 版本 10.12.1 (16B2555)
  • cuda_8.0.47_mac, cudnn-8.0-osx-x64-v5.1
  • TensorFlow 0.11.0rc2 通过以下方式安装：export TF_BINARY_URL=https://storage.googleapis.com/tensorflow/mac/gpu/tensorflow-0.11.0rc2-py2-none-any.whl

更新我发现选择 CPU 还是 GPU 实现取决于输入的大小。使用这个测试脚本：

from __future__ import print_function
import numpy as np
import tensorflow as tf
from time import time

NUM_VALUES = 1310720

input = np.arange(0, NUM_VALUES, dtype = float)

module = tf.load_op_library('custom_ops.so')
with tf.Session(config = tf.ConfigProto(log_device_placement = True)):
  start = time(); print(module.add_b(input, 8.).eval()); end = time(); print(end - start)

.. 当NUM_VALUES为 1310720 或更少时，则使用 CPU 实现。当NUM_VALUES为 1310721 或更多时，则使用 GPU 实现。

是否有 (1310720 * 8 bytes per double = ) 10 MiB 截止？如果是这样，我该如何覆盖它？AddB() 操作很简单，但对于更复杂的自定义操作，10 MiB 可能对于选择 GPU 实现来说太大了。

4

3 回答 3

3

我刚刚阅读了TensorFlow 问题 #2054 - 在 GPU 上手动放置具有 CPU 和 GPU 实现的自定义运算符将始终运行 CPU 版本，并且运行 CPU 实现的行为似乎是 TensorFlow 的一个功能，称为“常量折叠”。当 TensorFlow 在第一次运行之前优化图形时，涉及常量的操作通常在 CPU 上进行评估，因为认为 CPU 和 GPU 实现应该产生相同的结果。说得通。

禁用此行为的两种方法是：

  1. 禁用图形优化：

    from __future__ import print_function
import numpy as np
import tensorflow as tf
from time import time

NUM_VALUES = 10

input = np.arange(0, NUM_VALUES, dtype = float)

custom_ops_module = tf.load_op_library('custom_ops.so')

config = tf.ConfigProto(log_device_placement = True)
config.graph_options.optimizer_options.opt_level = -1

with tf.Session(config = config):
  start = time(); print(custom_ops_module.add_b(input, 8.).eval()); end = time(); print(end - start)

  2. 不使用常量，例如，将值输入占位符：

    from __future__ import print_function
import numpy as np
import tensorflow as tf
from time import time

NUM_VALUES = 10

custom_ops_module = tf.load_op_library('custom_ops.so')

graph = tf.Graph()
with graph.as_default():
  input = tf.placeholder(tf.float64, shape = (NUM_VALUES,))
  b = tf.placeholder(tf.float64, shape = ())
  result = custom_ops_module.add_b(input, b)

with tf.Session(graph = graph, config = tf.ConfigProto(log_device_placement = True)) as session:
  feed_dict = {
    input: np.arange(0, NUM_VALUES, dtype = float),
    b: 8.,
  }
  start = time(); print(session.run([result], feed_dict = feed_dict)); end = time(); print(end - start)
于 2016-11-24T16:07:01.100 回答
2

我认为模板实例化可能不正确：

template <typename Device, typename T>
class AddBOp : public OpKernel {
...
}

REGISTER_KERNEL_BUILDER(
    Name("AddB")
    .Device(DEVICE_CPU)
    .TypeConstraint<float>("T"),
    AddBOp<CPUDevice, float>);

接着：

template <typename T>
class AddBOp<GPUDevice, T> : public OpKernel {
...
}

REGISTER_KERNEL_BUILDER(
    Name("AddB")
    .Device(DEVICE_GPU)
    .TypeConstraint<float>("T"),
    AddBOp<GPUDevice, float>);

我认为为 GPU 注册 AddB 会实例化与第一个实现匹配的对象，而不是第二个（第一个实现有两个模板参数，第二个实现有一个）。

您可以通过在第二次注册中调用 AddBOp < float > 来解决此问题，但我建议使用更好的名称以避免混淆。

于 2016-11-23T00:17:05.760 回答
0

据此可能是由于内存碎片管理，尝试：

with tf.device('/gpu:0'):

或链接页面中用于调整内存碎片选项的片段。

编辑：要查看是否是这种情况，请尝试：

from __future__ import print_function
import numpy as np
import tensorflow as tf
from time import time

NUM_VALUES = 10

input = np.arange(0, NUM_VALUES, dtype = float)

custom_ops_module = tf.load_op_library('custom_ops.so')

config = tf.ConfigProto(log_device_placement = True)
config.gpu_options.allow_growth = True

with tf.Session(config = config):
    start = time(); print(custom_ops_module.add_b(input, 8.).eval()); end = time(); print(end - start)
于 2016-11-24T16:06:15.213 回答