Search code examples
tensorflowflatbufferstensorflow-lite

Understanding of SparsityParameters in tensorflow lite schema

I'm trying to understand the sparse tensor with tensorflow lite schema, which is very hard to do for me though.

Luckily, there is only one json example that be made from this schema(tensorflow/lite/testdata/sparse_tensor.json).

"sparsity": {
            "traversal_order": [0, 1, 2, 3], 
            "block_map": [0, 1], 
            "dim_metadata": [
              {   
                "format": "DENSE",
                "dense_size": 2
              },  
              {   
                "format": "SPARSE_CSR",
                "array_segments_type": "Uint8Vector",
                "array_segments": {"values": [0, 2, 3]},
                "array_indices_type": "Uint8Vector",
                "array_indices": {"values": [0, 1, 1]} 
              },  
              {   
                "format": "DENSE",
                "dense_size": 2
              },  
              {   
                "format": "DENSE",
                "dense_size": 2
              }   
            ]   
          }   
"buffers": [
    {   
    },  
    {
      "data": [
        1, 0, 0, 4,
        2, 3, 0, 0,
        5, 0, 0, 6
      ]
    }
  ]

And, this is the schema I refer to(tensorflow/lite/schema/schema.fbs).

table DimensionMetadata {
  // Whether a dimension is dense or sparse.
  format:DimensionType;
  // Index metadata used for a dimension.
  //   - If format is DimensionType.DENSE then we use the dense_size field to
  //     store the size of that dimension. Each index in that dimension is
  //     stored implicitly.
  //   - If format is DimensionType.SPARSE_CSR then we use array_segments and
  //     array_indices to encode that dimension. array_segments represents how
  //     to segment the indices array, each segment corresponds to one element
  //     in the previous dimension. array_indices represents the index of the
  //     non-zero elements within this dimension (as those in the CSR matrix
  //     format, where the first array is row pointers and the second array is
  //     column indices).
  dense_size:int;
  array_segments:SparseIndexVector;
  array_indices:SparseIndexVector;
}

// Parameters to encode a sparse TfLite tensor.
table SparsityParameters {
  // The traversal order of the dimensions defined in the `shape` field of the
  // conceptual dense tensor. For a n-dimensional tensors with dims (d0, d1,
  // ..., dn-1),
  //   - if not block sparse, the traversal_order is just a permutation of (d0,
  //     ..., dn-1). For example, a 2-D matrix stored in row-major order would
  //     have traversal_order = (d0, d1).
  //   - if block sparse with a k-dimensional block (0 <= k <= n), the
  //     traversal_order has n + k elements. The first n elements are still a
  //     permutation of (d0, ..., dn-1). The lask k elements are a permutation
  //     of (dn, ..., dn+k-1), defining how to traverse a block internally. For
  //     example, a 2-D matrix with 2-D blocks, both stored in row-major order
  //     would have traversal_order = (d0, d1, d2, d3).
  traversal_order:[int];
  // For an n-dimensional tensor with a k-dimensional block (0 <= k <= n),
  // stores how a block dimension in (dn, ..., dn+k-1) maps to the original
  // tensor dimension in (d0, ..., dn).
  // It's stored in the order of (dn, ..., dn+k-1).
  // If not block-sparse, this field is NULL.
  block_map:[int];
  // In the traversal order defined above, the metadata needed for
  // each dimension to locate the non-zero values in the original dense tensor.
  // The size of the dim_metadata array = the size of the traversal_order array
  // = n + k.
  dim_metadata:[DimensionMetadata];
}

As you can see above, there is a buffer that contains sparse tensor's contents.

"buffers": [
    {   
    },  
    {
      "data": [
        1, 0, 0, 4,
        2, 3, 0, 0,
        5, 0, 0, 6
      ]
    }
  ]

AFAIK, if I want to generate a sparse tensor like above, I have to write a code like below.

st1 = tf.compat.v1.sparse.SparseTensor(
    indices=[[0, 0], [0, 3], [1, 0], [1, 1], [2, 0], [2, 2]], values=buffers, dense_shape=[4, 4])

But, above json example don't match with my understanding at all.

  1. I think array_indices should be [0, 3, 0, 1, 0, 3] rather than [0, 1, 1] and array_segements be [0, 2, 2, 4, 4, 6] rather than [0, 2, 3].

Moreover, actually, none of that comments in schema understand at all..

  1. What does the metadata with DENSE format stand for?
{   
  "format": "DENSE",
  "dense_size": 2
},

As shcema's comment said, it is a field to store the size of that dimension.

But, which dimension has got value "2"? The shape is [4, 4]. I can't even infer where the number 2 come from.

  1. What is the "block"?

AFAIK, block is a box that contains non-zero value. But then, there are lots of blocks in above buffer I think.

If there is this block like,

1 2 0
3 4 0
0 0 0

I would say, it has one 2x2 block.

But there are six 1x1 blocks in above buffer.

Then, how can I make this things a block map..?

Actually, I don't know about traversal_order as well but if I knew the above things, I would understand it as well.

Please somebody help me..

Solution

  • Currently TFLite uses a different sparse tensor representation from Tensorflow. The format it uses is called TACO. Please refer to this paper for more details: http://tensor-compiler.org/kjolstad-oopsla17-tensor-compiler.pdf, section 3.

    The tf.SparseTensor is not meant to work with this, since it uses the COO format.

    A block is an inner sub-unit of a tensor that needs to be stored together. It can contain 0-valued element. The example flatbuffer shows a 2-D 4x4 tensor with a 2-D 2x2 inner block. TFLite uses blocked sparse tensor to take advantage of the NEON SIMD instructions.