tensorflow源码解析之framework-tensor

目录

  1. 什么是tensor
  2. tensor继承体系
  3. 与Eigen3库的关系
  4. 什么是tensor_reference
  5. tensor_shape
  6. tensor_slice
  7. 其它结构
  8. 关系图
  9. 涉及的文件
  10. 迭代记录

1. 什么是tensor

TF全称叫做TensorFlow,可见tensor的重要性。它本质上是一个对高维数据的封装,提供了丰富的API。在线性代数中,我们常用向量、矩阵来表示数据,而在深度学习应用中,有对更高维数据的需求。比如在对图像进行处理时,彩色图像本身就带有三维的信息(长、宽、颜色通道),通常还需要对彩色图像进行批处理,这样待处理的数据变为四维,在一些特殊的情形下,往往还需要更高维度的数据。如果针对每种多维数据定义一种结构,必然给计算带来不便。TF的做法是,为高维数据定义统一的类型Tensor。

但高维数据的概念有点抽象,为了让大家能对Tensor内部的数据结构有个直观的印像,我们先看一下Tensor类的私有数据成员:

class Tensor {
    //...
  private:
    TensorShape shape_;
    TensorBuffer* buffer_;
}

这两个结构都没有见过,不过没关系,只把它们当做张量的形状和底层数据指针就好了。Tensor作为一个核心数据类,必然提供了很多API,比如常规的构造、析构、赋值、复制、数值属性获取等。除此之外,还提供了两类比较特殊的接口,我们举例说明:

class Tensor {
  public:
    //...
    //与proto数据的相互转化
    bool FromProto(const TensorProto& other);
    void AsProtoField(TensorProto* proto);
    //为底层数据创建新视图
    template <typename T> typename TTypes<T>::Vec vec();
    template <typename T> typename TTypes<T>::Matrix matrix();
    template <typename T> typename TTypes<T, NDIMS>::Tensor tensor();
}

其中第一类将Tensor与序列化的proto之间相互转化,便于在设备之间传递Tensor。第二类是为当前的Tensor的底层数据提供另外一种视图,我们重点来说一下视图的概念。

回顾Tensor包含的私有数据,TensorBuffer* buffer_是一个指向底层数据的指针,关于它的结构在下文中会详细说明。这意味着,Tensor并不包含实际的底层数据,它实际上只是对底层数据的一种视图。同样一份底层数据,可以提供多种视图。比如对于一个长度为12的数组,如果把它看做向量,它是一个1x12的向量,如果把它看作矩阵,可以认为是3x4或者2x6的矩阵,如果把它当作张量,可以认为是3x2x2的张量。通过这种方法,我们可以对同一份底层数据进行复用,避免了重复申请内存空间,提升了效率。

graph TB A("Tensor A, shape=[3,4]")-->D(底层数据TensorBuffer) B("Tensor B, shape=[2,6]")-->D(底层数据TensorBuffer) C("Tensor C, shape=[3,2,2]")-->D(底层数据TensorBuffer)

顺便提一句,numpy中对多维数组的实现,也是同样的原理。

细心的读者可能发现了,在对底层数据创建新视图时,返回了一种奇怪的数据类型typename TTypes<T>::Vec,这涉及TF中的Tensor与Eigen3库的关系,我们将在下文中详细说明。

2. tensor继承体系

接下来我们看一下TensorBuffer到底是什么样的结构。它只是一个继承自引用计数类的虚拟接口,不包含任何实现:

class TensorBuffer : public core::RefCounted {
    //...
}

因此怀疑,TensorBuffer只是一个提供接口的基类,实际上能用的只是它的子类。我们看下它的继承结构:

class BufferBase : public TensorBuffer {
    //...
}
class Buffer : public BufferBase {
    //...
  private:
    T* data_;
    int64 elem_;
}

结构已经非常清晰了,BufferBase类继承自TensorBuffer,它除了包含一个内存分配器指针外,还对基类中的部分API进行了实现。而Buffer类是实际可用的,它包含了指向实际数据的指针data_以及元素数量elem_。

另外还要说明一点,Buffer除了申请内存之外,还能调用目标类的构造和析构函数,初始化Buffer的内容,TF为此设计了很多辅助类和函数,这里就不一一赘述了。

Tensor的继承体系图如下:

graph TB A(core::RefCounted)-->|派生|B(TensorBuffer) B(TensorBuffer)-->|派生|C(BufferBase) C(BufferBase)-->|派生|D(Buffer) C(BufferBase)-.包含.->E(Allocator* alloc_) D(Buffer)-.包含.->F(T* data_) D(Buffer)-.包含.->G(int64 elem_) H(Tensor)-.包含.->B(TensorBuffer)

3. 与Eigen3库的关系

刚才提到了,当为Tensor的数据提供不同视图的时候,返回了一种奇怪的数据TTypes<T>::Vec,这种数据为TF中的Tensor和Eigen3库中的Tensor建立了联系。我们在tensor_types.h文件中,找到了这种类型的定义:

struct TTypes {
    typedef Eigen::TensorMap<Eigen::Tensor<T,NDIMS,Eigen::RowMajor,IndexType>,Eigen::Aligned> Tensor;
    typedef Eigen::TensorMap<Eigen::Tensor<T,1,Eigen::RowMajor,IndexType>,Eigen::Aligned> Vec;
    //...
}

原来,对Eigen3库中Tensor的使用在这里。由于这种定义被包裹在TTypes结构体中,所以不会与外部TF自定义的Tensor造成冲突。

重新回到Tensor的定义,我们发现,原来在对Tensor底层数据提供多种视图的时候,返回的已经不是Tensor结构,而是TTypes::TensorMap,这是否意味着,TF中定义的Tensor只是对Eigen::Tensor的一种封装呢?我们追根溯源,找到vec函数的实现:

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

template <typename T, size_t NDIMS>
typename TTypes<T, NDIMS>::Tensor Tensor::tensor() {
    CheckTypeAndIsAligned(DataTypeToEnum<T>::v());
    return typename TTypes<T, NDIMS>::Tensor(base<T>(), shape().AsEigenDSizes<NDIMS>());
}

跟我们预想的完全一样,在对vec函数的调用中,调用了tensor函数,而这个函数的作用,就是将TF中定义的Tensor转变为TTypes::Tensor,而后者就是Eigen::TensorMap,也就是说,tensor返回的本质上是一个Eigen::TensorMap。另外,我们知道base()和shape()两个函数,分别返回了TensorBuffer指针和TensorShape,因此实际上就是使用TF中Tensor存储的数据,作为了Eigen::TensorMap的构造函数的参数。

可以说,TF中的Tensor实际上是对Eigen::TensorMap的一种高级封装,它不是简单的在私有数据成员包含后者,而是包含了构造后者所需要的数据,在需要后者的时候,构造并返回。这种方式,使得TF中的Tensor既能利用Eigen高效的张量计算方法,也能为Tensor定制一些API。

4. 什么是tensor_reference

Tensor类的对象除了包含指向底层数据的指针外,还包含了对数据形状和类型的描述(通过TensorShape),如果我们并不关心这些,直接使用Tensor会增加构建或者移动的负担。因此TF推出了tensor_reference这个类,它仅包含了一个指向TensorBuffer的指针,并且每增加一个TensorReference对象,就会增加一个针对底层TensorBuffer的引用计数。因此针对TensorReference来说,我们唯一能做的就是在用完之后Unref掉,否则会造成内存泄漏。

class TensorReference {
  public:
    //...
  private:
    TensorBuffer* buf_;
}

5. tensor_shape

TensorShape显然包含的是张量形状相关的信息,但其实不仅如此,它还包含了对张量数据类型的描述。TensorShape相关的核心类继承体系如下:

graph LR I(TensorShapeRep)-->|派生|J(TensorShapeBase) J(TensorShapeBase)-->|派生|K(TensorShape) J(TensorShapeBase)-->|派生|L(PartialTensorShape)

首先来看一下,最底层的TensorShapeRep的私有数据成员:

class TensorShapeRep {
    //...
  private:
    union {
        uint8 buf[16];
        Rep64* unused_aligner;//除了强制u_与指针对齐外,没有任何作用
    } u_;
    int64 num_elements_;
}

buf这个数组很有意思,它的前12个元素用来存储形状,虽然Tensor最高能支持到256维的张量,但最常用的不超过3维,为了效率,TF提供了三种利用这12个字节的方式,如下:

struct Rep16 {
    uint16 dims_[6];//最多可表示6维的张量,每一维的长度不超过2^16-1
};
struct Rep32 {
    uint32 dims_[3];//最多可表示3维的张量,每一维的长度不超过2^32-1
};
struct Rep64 {
    gtl::InlinedVector<int64, 4>* dims_;//支持任意维度的张量
};

剩下的4个字节也不能浪费,在第14-16个字节中,分别存储了张量中的数据类型编号、张量的维度数目、张量维度的表示类型(Rep16, Rep32, Rep64)。由于张量维度的数目是用一个字节存储的,因此最多支持256维。可惜笔者目前仍没有发现第13个字节的作用,有发现的读者欢迎告知我。

TensorShapeBase类并没有添加额外的数据成员,它只是添加了一些允许我们修改张量维度的API接口。而TensorShape类也只是添加了一些对形状进行检查和比较的接口,没有新增数据成员。

最后再来看下PartialTensorShape类,在构造一个张量的形状时,如果对于某些维度我们还不知道具体的维度值,可以把这个维度设为未知,因此就会用到PartialTensorShape类,这个类中也包含了一些未知维度操作的API,这里就不详述了。

6. tensor_slice

TensorSlice类表示一个张量的索引,它的数据结构非常简单:

class TensorSlice {
    //...
  private:
    gtl::InlinedVector<int64,4> starts_;
    gtl::InlinedVector<int64,4> lengths_;
}

分别是每一个维度索引的开始位置和索引长度,由此我们也知道,TF对Tensor只支持连续索引,不支持间隔索引。
由于TensorSlice用途广泛,对其进行初始化的方法也多种多样,包括:

  • 创建空索引
  • 从单个维度创建(当创建全索引时)
  • 从一个整数对数组创建
  • 从一个TensorSliceProto创建
  • 从一个字符串描述中创建

7. 其它结构

为了方便对张量和与之相关的数据结构进行序列化,TF设计了很多protos,理解起来相对简单,现只说明下它们的用途,感兴趣的读者可以去看源代码。

message TensorDescription;//张量的描述,包括数据类型、形状、内存分配信息
message TensorProto;//张量的数据类型,版本,原始数据等
message VariantTensorDataProto;//对DT_VARIANT类型的序列化表示
message TensorShapeProto;//张量形状
message TensorSliceProto;//张量索引

8. 关系图

graph TB A(core::RefCounted)-->|派生|B(TensorBuffer) B(TensorBuffer)-->|派生|C(BufferBase) C(BufferBase)-->|派生|D(Buffer) C(BufferBase)-.包含.->E(Allocator* alloc_) D(Buffer)-.包含.->F(T* data_) D(Buffer)-.包含.->G(int64 elem_) H(Tensor)-.包含.->B(TensorBuffer) I(TensorShapeRep)-->|派生|J(TensorShapeBase) J(TensorShapeBase)-->|派生|K(TensorShape) J(TensorShapeBase)-->|派生|L(PartialTensorShape) H(Tensor)-.索引结构.->M(TensorSlice) H(Tensor)-.形状和数据类型描述.->K(TensorShape) H(Tensor)-.转换为.->N(TTypes::Tensor) N(TTypes::Tensor)-.转换为.->O(Eigen::TensorMap)

9. 涉及的文件

  • tensor
  • tensor_reference
  • tensor_types
  • tensor_shape
  • tensor_slice
  • tensor_description

10. 迭代记录

  • v1.0 2018-08-26 文档创建
  • v2.0 2018-09-09 文档重构

posted on 2018-08-26 14:26  jicanghai  阅读(3493)  评论(0编辑  收藏  举报

导航