Press "Enter" to skip to content

Tensorflow 的 Tensor 和 OpKernel 分析

本站内容均来自兴趣收集,如不慎侵害的您的相关权益,请留言告知,我们将尽快删除.谢谢.

的用户 API 可以参考这里,这里做一下简单介绍。Tensor是各种维度的向量和矩阵的统称,分 Tensor 和 SparseTensor。和 Tensor 不同,SparseTensor 存的是值以及值对应的 index,而 Tensor 存的是完整的矩阵。

举个例子。

import  as tfa = tf.constant([1, 1])b = tf.constant([2, 2])c = tf.add(a, b)sess = tf.InteractiveSession()print("a[0]=%s, a[1]=%s" % (a[0].eval(), a[1].eval()))print("c = %s" % c.eval())sess.close()

输出对应如下。

a[0]=1, a[1]=1c = [3 3]

Tensor 只有 eval 以后才能获得结果,是懒计算的。

Tensor 的实现

Tensor (tensorflow/tensorflow/core/framework/tensor.h) 依赖 TensorShape(tensorflow/tensorflow/tensorflow/core/framework/tensor_shape.h) 和 TensorBuffer (tensorflow/tensorflow/core/framework/tensor.h) 两个成员。

TensorShape 主要负责记录张量的形状。

TensorBuffer主要负责管理Tensor的内存,TensorBuffer继承自RefCounted(tensorflow/tensorflow/core/lib/core/refcount.h),具有引用计数的功能,用于对内存进行管理。

// Interface to access the raw ref-counted data buffer.class TensorBuffer : public core::RefCounted { public:  ~TensorBuffer() override {}  ...

他们的对应的关系如下。

Tensor 从下往上看,其实就是一个带”形状“的内存,和 NumPy 的数组是差不多的。

对于一个线性回归来说,是最简单也最好理解的模型,方便分析底层的代码实现。

$$ Y=XW+b $$

损失函数用平方差定义的,优化器是提督下降,这样一个模型可以用一下的 Python 代码实现,这个代码片段是截取的,如果要完整运行这个例子可以在这里复现。

import tensorflow as tfsess = tf.InteractiveSession()x = tf.placeholder(tf.float32,[None, x_train.shape[1]])y = tf.placeholder(tf.float32,[None, 1])w = tf.Variable(tf.zeros([x_train.shape[1],1]))b = tf.Variable(tf.zeros([1])) # placeholder 不用加 Nonepred = tf.add(tf.matmul(x, w), b)init = tf.global_variables_initializer()cost = tf.reduce_mean(tf.square(y - pred))optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001).minimize(cost)epochs = 3000init.run()for epoch in range(0, epochs):    optimizer.run(feed_dict={x:x_train, y:y_train})    c = cost.eval(feed_dict = {x:x_train,y:y_train})    if epoch%100 == 0:        print_percentage(int(c*100))print('nEpoch: {0}, Error: {1}'.format(epoch+1, c))b_value = b.eval()w_value = w.eval()# Predicted Labelsy_pred = pred.eval(feed_dict={x: x_test})# Mean Squared Errormse = tf.reduce_mean(tf.square(y_pred - y_test))print("MSE", mse.eval())sess.close()

对应的训练结果。

Epoch: 3000, Error: 0.25882452726364136MSE 0.30620116674806463

可以看到,线性回归的模型主要依赖的两个 Operation 分别是tf.addtf.matmul,其他的复杂模型也是类似的逻辑,对应的OpKernel分别是AddOpMatMulOp,这里可以看一下具体的实现。

如果有源代码,可以连着源代码用 bazel 编译,可以参照这里自己编写一个 Op。

MatMulOp的实现在/tensorflow/core/kernels/matmul_op.cc下面,定义在tensorflow/tensorflow/core/ops/math_ops.cc下面。AddOp/tensorflow/core/kernels/matmul_op.cc,实现在/tensorflow/core/kernels/cwise_op_add_1.cc下面,依赖/tensorflow/tensorflow/core/kernels/cwise_ops_common.hcommon的定义。

Add用的是Eigenadd/tensorflow/tensorflow/core/kernels/cwise_ops.h,依赖third_party/Eigen/src/Core/functors/BinaryFunctors.h

举个例子,看一下MatMulOpMatMulOp的构造函数里面有一个OpKernelConstruction可以初始化OpKernel,通过OpKernel可以获得这个Op的参数比如transpose_a等等。

template <typename Device, typename T, bool USE_CUBLAS>class MatMulOp : public OpKernel { public:  explicit MatMulOp(OpKernelConstruction* ctx)      : OpKernel(ctx), algorithms_set_already_(false) { // 在执行构造函数之前,执行两个成员的构造函数    OP_REQUIRES_OK(ctx, ctx->GetAttr("transpose_a", &transpose_a_));    OP_REQUIRES_OK(ctx, ctx->GetAttr("transpose_b", &transpose_b_));    LaunchMatMul<Device, T, USE_CUBLAS>::GetBlasGemmAlgorithm(        ctx, &algorithms_, &algorithms_set_already_);    use_autotune_ = MatmulAutotuneEnable();  }

每个 OpKernel 都要实现一个 Compute 函数,可以看到这个 Compute 函数首先检查了两个 Tensor 是否是矩阵,然后检查两个矩阵的形状是否符合矩阵相乘的条件,然后根据形状分配TensorShape并且根据TensorShape分配新的Tensor(其实顺便分配的TensorBuffer的内存空间)。然后通过LaunchMatMul真正执行相乘操作,因为这个计算过程,可能是用了 GPU,所以模版是带 Device 的(GPUDevice/CPUDevice)。

void Compute(OpKernelContext* ctx) override {  const Tensor& a = ctx->input(0);  const Tensor& b = ctx->input(1);  // Check that the dimensions of the two matrices are valid.  OP_REQUIRES(ctx, TensorShapeUtils::IsMatrix(a.shape()),              errors::InvalidArgument("In[0] is not a matrix"));  OP_REQUIRES(ctx, TensorShapeUtils::IsMatrix(b.shape()),              errors::InvalidArgument("In[1] is not a matrix"));  Eigen::array<Eigen::IndexPair<Eigen::DenseIndex>, 1> dim_pair;  dim_pair[0].first = transpose_a_ ? 0 : 1;  dim_pair[0].second = transpose_b_ ? 1 : 0;  OP_REQUIRES(      ctx, a.dim_size(dim_pair[0].first) == b.dim_size(dim_pair[0].second),      errors::InvalidArgument(          "Matrix size-incompatible: In[0]: ", a.shape().DebugString(),          ", In[1]: ", b.shape().DebugString()));  int a_dim_remaining = 1 - dim_pair[0].first;  int b_dim_remaining = 1 - dim_pair[0].second;  TensorShape out_shape(      {a.dim_size(a_dim_remaining), b.dim_size(b_dim_remaining)});  Tensor* out = nullptr;  OP_REQUIRES_OK(ctx, ctx->allocate_output(0, out_shape, &out));  if (out->NumElements() == 0) {    // If a has shape [0, x] or b has shape [x, 0], the output shape    // is a 0-element matrix, so there is nothing to do.    return;  }  if (a.NumElements() == 0 || b.NumElements() == 0) {    // If a has shape [x, 0] and b has shape [0, y], the    // output shape is [x, y] where x and y are non-zero, so we fill    // the output with zeros.    functor::SetZeroFunctor<Device, T> f;    f(ctx->eigen_device<Device>(), out->flat<T>());    return;  }  LaunchMatMul<Device, T, USE_CUBLAS>::launch(      ctx, a, b, dim_pair, &algorithms_, use_autotune_, out);}

LaunchMatMul继承自LaunchMatMulBase,在LaunchMatMulBase当中调用了functor::MatMulFunctor,这个 functor 主要就会执行乘法操作,在这之前会检查一下是否其中一个元素是 vector,这样可以直接优化算出来,而不用 Eigen 库来算,这样更快,这个目前看到的是 CPU 的路径。

template <typename Device, typename T>struct LaunchMatMulBase {#ifGOOGLE_CUDA  typedef se::blas::AlgorithmType AlgorithmType;#else  typedef int64 AlgorithmType;#endif// GOOGLE_CUDA  static void launch(
OpKernelContext* ctx, const Tensor& a, const Tensor& b,
const Eigen::array<Eigen::IndexPair<Eigen::DenseIndex>, 1>& dim_pair,
std::vector<AlgorithmType>* algorithms, bool use_aututone, Tensor* out){#ifndefTENSORFLOW_USE_SYCL // An explicit vector-matrix multiply is much better optimized than an // implicit one and this is a bottleneck during non-batched inference. bool was_vector = ExplicitVectorMatrixOptimization<T>(a, b, dim_pair, out); if (!was_vector) {#endif// TENSORFLOW_USE_SYCL functor::MatMulFunctor<Device, T>()(ctx->eigen_device<Device>(), out->matrix<T>(), a.matrix<T>(), b.matrix<T>(), dim_pair);#ifndefTENSORFLOW_USE_SYCL }#endif// TENSORFLOW_USE_SYCL } static void GetBlasGemmAlgorithm(OpKernelConstruction* ctx,
std::vector<int64>* algorithms,
bool* algorithm_set_flag){}};

MatMulFunctor在设备d上计算矩阵相乘的结果,其中调用的是MatMul<CPUDevice>

template <typename Device, typename In0, typename In1, typename Out,          typename DimPair>void MatMul(const Device& d, Out out, In0 in0, In1 in1,
const DimPair& dim_pair){ out.device(d) = in0.contract(in1, dim_pair);}

这里的contract调用的是TensorContractionOp(third_party/unsupported/Eigen/CXX11/src/Tensor/TensorContraction.h),跟之前说的一样,这个Op是计算图的一部分,要通过eval来做计算,计算结果是eval驱动的。TensorContractionOp的构造函数就是,这负责构建左表达式和右表达式。

EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorContractionOp(    const LhsXprType& lhs, const RhsXprType& rhs, const Indices& dims)    : m_lhs_xpr(lhs), m_rhs_xpr(rhs), m_indices(dims) {}

真正的计算过程在TensorContractionEvaluatorBase里面,真正执行计算过程,计算细节就省略了主要是矩阵相乘。

CUDA

如果条件编译GOOGLE_CUDA的话,会使用 GPU 的代码,对应会调用到 steam executor,这个以后具体分析。

总结

Tensorflow 基于图模型的,并且是懒计算的,通过扩展可以自己用 C++ 实现新的 Op,并且也可以观察默认自带的 OpKernel 是如何实现的,对于理解 Tensorflow 的工作流程会有很大的帮助。Tensorflow 本身依赖了 Eigen,CUDA 等线性代数库或者 GPU 计算库,要看懂代码还是要多学一点线代的知识,比如Contraction这个概念我也是第一次晓得。

Be First to Comment

发表评论

电子邮件地址不会被公开。 必填项已用*标注