自定义张量流C ++ op中的特征和，在GPUDevice上具有Eigen :: half

Question

我正在尝试创建自定义张量流Op。我已经使用https://www.tensorflow.org/versions/master/how_tos/adding_an_op/index.html和普通C ++获得了一些操作。

使用Eigen C ++框架及其.sum缩减器时会出现问题。它适用于具有double，float和Eigen::half类型的CPU，但在GPU上，它在编译时使用Eigen::half时会中断。

我已将问题缩减为l2loss_op的{​​{1}}副本，我已将L2Loss重命名为CustomL2Loss（否则会出现名称冲突）。请参阅：https://github.com/tensorflow/tensorflow/tree/master/tensorflow/core/kernels

错误很长：https://gist.github.com/AndreasMadsen/4335215cd4293daad3cad745bbeae82a

因为它适用于除<GPUDevice, Eigen::half>之外的所有其他情况（通过在l2loss_op.cu.cc中评论该行确认），我正在考虑这是否是一个张量流问题。但我可以自己编译tensorflow。

Answer 1

支持半浮动需要cuda架构大于或等于3.5。您需要使用-arch compute_35标志进行编译以启用相应的指令。