如何将Eigen::TensorMap转换为float*？TensorFlow自定义Op开发遇阻

在实现TensorFlow r0.12自定义Op时，要调用以float指针为参数的forward()函数，核心是正确获取张量底层的float类型数据指针。下面是具体的实现步骤和常见问题的解决方案：

1. 先确保输入张量的类型匹配 #

首先在Op注册阶段，必须明确指定输入张量的类型为float（对应TensorFlow的DT_FLOAT），避免类型不匹配的问题：

REGISTER_OP("YourCustomOp")
    .Input("input_tensor: float")
    .Output("output_tensor: float");

2. 在Kernel的Compute方法中获取float指针 #

继承OpKernel的自定义Kernel类，在Compute方法里按以下步骤获取指针：

void Compute(OpKernelContext* context) override {
    // 第一步：获取输入张量
    const Tensor& input_tensor = context->input(0);

    // 第二步：校验张量类型，避免类型错误
    OP_REQUIRES(context, input_tensor.dtype() == DT_FLOAT,
                errors::InvalidArgument("输入必须是float类型，当前类型为: ", input_tensor.dtype()));

    // 第三步：获取flat视图的float指针
    const float* input_data_ptr = input_tensor.flat<float>().data();

    // 现在可以直接调用你的forward函数
    // forward(input_data_ptr, ...);

    // 后续处理输出张量的逻辑...
}

3. 常见报错的排查与解决 #

如果还是出现转换失败的报错，大概率是以下场景导致的：

• 张量运行在GPU上：如果你的Op绑定了GPU设备，直接获取CPU指针会出错。需要先把张量拷贝到CPU，或者在GPU环境下处理数据。如果必须用CPU指针，可以这样做：

  Tensor cpu_tensor;
  OP_REQUIRES_OK(context, context->device()->MakeTensorFromProto(
      input_tensor.AsProtoTensorContent(), DEVICE_CPU, &cpu_tensor));
  const float* input_data_ptr = cpu_tensor.flat<float>().data();
  

• 输入是多维张量：flat<float>()会自动把多维张量展平为一维连续视图，此时获取的指针依然有效（TensorFlow张量默认是连续行优先存储）。如果需要保留维度信息，也可以用多维张量视图来取指针：

  // 以二维张量为例
  auto input_matrix = input_tensor.tensor<float, 2>();
  const float* input_data_ptr = &input_matrix(0, 0); // 获取第一个元素的指针
  

• 需要修改张量数据：如果forward()需要修改数据，不能用const float*，需要获取可写指针。这时候可以创建临时张量或者直接操作输出张量：

  Tensor* temp_tensor = nullptr;
  OP_REQUIRES_OK(context, context->allocate_temp(DT_FLOAT, input_tensor.shape(), &temp_tensor));
  float* mutable_data_ptr = temp_tensor->flat<float>().data();
  // forward(mutable_data_ptr, ...);
  

4. 验证Op的正确性 #

按照官方文档的验证步骤，用Python测试你的Op：

import tensorflow as tf

# 加载编译好的Op库
custom_op_lib = tf.load_op_library('./your_custom_op.so')

# 创建测试输入
test_input = tf.constant([1.0, 2.0, 3.0], dtype=tf.float32)
test_output = custom_op_lib.your_custom_op(test_input)

# 运行会话验证结果
with tf.Session() as sess:
    print("输出结果:", sess.run(test_output))

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