TensorFlow源碼分析——Tensor與Eigen

TensorFlow底層操作的數據結構是Tensor（張量），可以表示多維的數據，其實現在core/framework/tensor.h中，對於tensor的理解主要分兩大塊：

1.Tensor的組成成分

2.Tensor是如何進行數學運算的（TensorFlow本質就是處理大量訓練數據集，在底層要實現深度學習常用的算法，自然涉及到一些數學運算，例如：Tensor相加、相減，softmax，reduction等操作）

 

Tensor的組成：

Tensor的組成可以分成3部分：DataType、TensorShape、TensorBuf. 即數據類型、張量的形狀和存儲數據的內存.

DataType：Tensor數據的類型，float、double、int等

TensorShape: Tensor的形狀. 這里要表示維度需存儲兩種元素，第一種是維數、第二種是每個維度的大小. 例如：一個矩陣的大小為3×4，那么其維數為2，第一個維度大小是3，第二個是4. 

TensorBuf:就是存儲數據的內存地址，比如一個Tensor是3×4的矩陣，並且類型時float，那么其TensorBuf就float*類型，並且長度是12，在實現上是基於模板實現的.

 

Tensor的運算：

Tensor的運算主要由kernels完成，每個kernel都完成不同的運算。其中CPU版本的kernel大部分都使用Eigen庫來實現，GPU版本的kernel少部分使用Eigen，大部分使用cuda編程實現。

Eigen unsupported模塊提供了同樣叫Tensor的類，主要就是能完成各種復雜的數學計算，並且是並行實現的（主要使用線程池或cuda），效率很高. 其API文檔可參考 https://github.com/PX4/eigen/blob/master/unsupported/Eigen/CXX11/src/Tensor/README.md. TensorFlow使用Eigen進行計算時，首先要將其自己實現的Tensor類的對象轉換成Eigen支持的Tensor對象. 可以再core/framework/tensor_types.h下面看到Eigen支持的各種數據類型，主要是TensorMap類. 並進行了重命名，其源碼如下所示：

struct TTypes {
  // Rank-<NDIMS> tensor of scalar type T.
  typedef Eigen::TensorMap<Eigen::Tensor<T, NDIMS, Eigen::RowMajor, IndexType>,
                           Eigen::Aligned>
      Tensor;
  typedef Eigen::TensorMap<
      Eigen::Tensor<const T, NDIMS, Eigen::RowMajor, IndexType>, Eigen::Aligned>
      ConstTensor;

  // Unaligned Rank-<NDIMS> tensor of scalar type T.
  typedef Eigen::TensorMap<Eigen::Tensor<T, NDIMS, Eigen::RowMajor, IndexType> >
      UnalignedTensor;
  typedef Eigen::TensorMap<
      Eigen::Tensor<const T, NDIMS, Eigen::RowMajor, IndexType> >
      UnalignedConstTensor;

  typedef Eigen::TensorMap<Eigen::Tensor<T, NDIMS, Eigen::RowMajor, int>,
                           Eigen::Aligned>
      Tensor32Bit;

  // Scalar tensor (implemented as a rank-0 tensor) of scalar type T.
  typedef Eigen::TensorMap<
      Eigen::TensorFixedSize<T, Eigen::Sizes<>, Eigen::RowMajor, IndexType>,
      Eigen::Aligned>
      Scalar;
  typedef Eigen::TensorMap<Eigen::TensorFixedSize<const T, Eigen::Sizes<>,
                                                  Eigen::RowMajor, IndexType>,
                           Eigen::Aligned>
      ConstScalar;

  // Unaligned Scalar tensor of scalar type T.
  typedef Eigen::TensorMap<
      Eigen::TensorFixedSize<T, Eigen::Sizes<>, Eigen::RowMajor, IndexType> >
      UnalignedScalar;
  typedef Eigen::TensorMap<Eigen::TensorFixedSize<const T, Eigen::Sizes<>,
                                                  Eigen::RowMajor, IndexType> >
      UnalignedConstScalar;

  // Rank-1 tensor (vector) of scalar type T.
  typedef Eigen::TensorMap<Eigen::Tensor<T, 1, Eigen::RowMajor, IndexType>,
                           Eigen::Aligned>
      Flat;
  typedef Eigen::TensorMap<
      Eigen::Tensor<const T, 1, Eigen::RowMajor, IndexType>, Eigen::Aligned>
      ConstFlat;
  typedef Eigen::TensorMap<Eigen::Tensor<T, 1, Eigen::RowMajor, IndexType>,
                           Eigen::Aligned>
      Vec;
  typedef Eigen::TensorMap<
      Eigen::Tensor<const T, 1, Eigen::RowMajor, IndexType>, Eigen::Aligned>
      ConstVec;

  // Unaligned Rank-1 tensor (vector) of scalar type T.
  typedef Eigen::TensorMap<Eigen::Tensor<T, 1, Eigen::RowMajor, IndexType> >
      UnalignedFlat;
  typedef Eigen::TensorMap<
      Eigen::Tensor<const T, 1, Eigen::RowMajor, IndexType> >
      UnalignedConstFlat;
  typedef Eigen::TensorMap<Eigen::Tensor<T, 1, Eigen::RowMajor, IndexType> >
      UnalignedVec;
  typedef Eigen::TensorMap<
      Eigen::Tensor<const T, 1, Eigen::RowMajor, IndexType> >
      UnalignedConstVec;

  // Rank-2 tensor (matrix) of scalar type T.
  typedef Eigen::TensorMap<Eigen::Tensor<T, 2, Eigen::RowMajor, IndexType>,
                           Eigen::Aligned>
      Matrix;
  typedef Eigen::TensorMap<
      Eigen::Tensor<const T, 2, Eigen::RowMajor, IndexType>, Eigen::Aligned>
      ConstMatrix;

  // Unaligned Rank-2 tensor (matrix) of scalar type T.
  typedef Eigen::TensorMap<Eigen::Tensor<T, 2, Eigen::RowMajor, IndexType> >
      UnalignedMatrix;
  typedef Eigen::TensorMap<
      Eigen::Tensor<const T, 2, Eigen::RowMajor, IndexType> >
      UnalignedConstMatrix;
};

 

在tensor.h可以看到一些轉換函數，比如：flat(拉伸)、vec(向量化)、matrix(矩陣化)等。部分源碼如下所示：

template <typename T>
  typename TTypes<T>::Vec vec() {
    return tensor<T, 1>();
  }

  template <typename T>
  typename TTypes<T>::Matrix matrix() {
    return tensor<T, 2>();
  }

  template <typename T, size_t NDIMS>
  typename TTypes<T, NDIMS>::Tensor tensor();

template <typename T>
  typename TTypes<T>::Flat flat() {
    return shaped<T, 1>({NumElements()});
  }

調用這些函數就能將tensor轉換成Eigen支持的類型，就可以直接使用Eigen提供的API進行計算. 像矩陣乘、二元操作、一元操作等都是通過調用Eigen實現的. 所以只要了解了Eigen的基本數據類型，就能輕松看懂tensor這部分源碼. 了解了Eigen操作基本類型的API就能看懂kernels的代碼. 

注：這里為什么不直接用Eigen::Tensor？原因是Eigen::Tensor只提供了最基本的數學運算，而有時候我們在運算的時候要進行不同程度的預處理，可以參考源碼中的cwise_ops_common.h中的

BinaryOp class以及reduction_ops_common.h中的ReductionOp class，這兩個kernels都在運算前進行了預處理. 如果沒有把Eigen::Tensor包裝成我們自己的Tensor，我們就需要調用大量的Eigen的API，代碼更難理解，這樣包裝之后，我們就只需要調用Eigen計算的一些API，而不用考慮獲取一些屬性. 另一個原因就是Tensor有一些屬性和方法，用Eigen::Tensor的API並不好實現，並且Eigen的Tensor在聲明的時候，維度是模板參數，因此只能為常量，這樣就帶來很多不便之處.  總之多包裝一層代碼層次感就更強，閱讀起來就更容易一些.