时间戳

首页 > 解决方案 > TensorFlow 算子源代码

tensorflow - TensorFlow 算子源代码

问题描述

我正在尝试查找 TensorFlow 的低级线性代数和矩阵算术运算符的源代码，以便在 CPU 上执行。例如，在tf.add()CPU 上执行的实际实现在哪里？据我所知，大多数线性代数运算符实际上是由 Eigen 实现的，但我想知道具体调用了哪些 Eigen 函数。

我尝试从高级 API 进行追溯，但这很困难，因为在将运算符放置在图形上与 TF 运行时实际执行运算符之间有很多步骤。

标签： tensorflow

解决方案

该实现隐藏在一些元模板编程之后（对于 Eigen 来说并不罕见）。

TensorFlow 中的每个操作都在某个时间点注册。在这里和这里Add注册。

REGISTER3(BinaryOp, GPU, "Add", functor::add, float, Eigen::half, double);

Operations 的实际实现是基于OpKernel. 该Add操作在BinaryOp::Compute类层次结构中实现BinaryOp : BinaryOpShared : OpKernel

在添加两个标量的情况下，整个实现只是：

functor::BinaryFunctor<Device, Functor, 1>().Right(
            eigen_device, out_flat, in0.template flat<Tin>(),
            in1.template scalar<Tin>(), error_ptr);

in0, in1传入的张量标量在哪里，Device是 GPU 还是 CPU，Functor是操作本身。其他行仅用于执行广播。

在此文件中向下滚动并展开REGISTER3宏说明如何将参数从 传递REGISTER3到functor::BinaryFunctor<Device, Functor, ...>。

您不能期望看到一些循环，因为 Eigen 使用表达式进行惰性求值和混叠。Eigen-“呼叫”在这里：

https://github.com/tensorflow/tensorflow/blob/7a0def60d45c1841a4e79a0ddf6aa9d50bf551ac/tensorflow/core/kernels/cwise_ops.h#L693-L696

推荐阅读