Search code examples
pythontensorflowtensorflow-datasets

Tensorflow: RaggedTensor.from_tensor flattening values from all arrays into one array, instead of preserving original number of arrays

In the official documentation, RaggedTensor.from_tensor will work something like this. 

x = [[1, 3, -1, -1], [2, -1, -1, -1], [4, 5, 8, 9]]
print(tf.RaggedTensor.from_tensor(x, padding=-1))

Output:

 <tf.RaggedTensor [[1, 3], [2], [4, 5, 8, 9]]>

Preserving the original number of arrays.

However, when working with a batch outputted by the dataset api iterator, it flattens it to one array. Here is the key parts of the code. 

dataset = dataset.padded_batch(3, padded_shapes=([None],[None]), padding_values=(tf.constant(-1, dtype=tf.int64)
                                                 ,tf.constant(-1, dtype=tf.int64)))
iterator = dataset.make_one_shot_iterator()
i, data = iterator.get_next()

data2= tf.RaggedTensor.from_tensor(data, padding=-1)

with tf.Session() as sess:
    print(sess.run([ data, data2 ]))
    print(sess.run([ data, data2 ]))
    print(sess.run([ data, data2 ]))

Here is the output

[array([[ 0,  1,  2,  3, -1],
       [ 2,  3,  4, -1, -1],
       [ 3,  6,  5,  4,  3]]), tf.RaggedTensorValue(values=array([0, 1, 2, 3, 2, 3, 4, 3, 6, 5, 4, 3]), row_splits=array([ 0,  4,  7, 12]))]
[array([[ 3,  9, -1, -1],
       [ 0,  1,  2,  3],
       [ 2,  3,  4, -1]]), tf.RaggedTensorValue(values=array([3, 9, 0, 1, 2, 3, 2, 3, 4]), row_splits=array([0, 2, 6, 9]))]
[array([[ 3,  6,  5,  4,  3],
       [ 3,  9, -1, -1, -1],
       [ 0,  1,  2,  3, -1]]), tf.RaggedTensorValue(values=array([3, 6, 5, 4, 3, 3, 9, 0, 1, 2, 3]), row_splits=array([ 0,  5,  7, 11]))]

Here is the full code to the minimal example to reproduce the results

!pip install -q tf-nightly
import math
import numpy as np
import tensorflow as tf

#Generate Test data
cells = np.array([[0,1,2,3], [2,3,4], [3,6,5,4,3], [3,9]])
mells = np.array([[0], [2], [3], [9]])
print(cells)

#Write test data to tf.records file
writer = tf.python_io.TFRecordWriter('test.tfrecords')
for index in range(mells.shape[0]):
    example = tf.train.Example(features=tf.train.Features(feature={
        'num_value':tf.train.Feature(int64_list=tf.train.Int64List(value=mells[index])),
        'list_value':tf.train.Feature(int64_list=tf.train.Int64List(value=cells[index]))
    }))
    writer.write(example.SerializeToString())
writer.close()

#Open tfrecords file and generate batch from data 
filenames = ["test.tfrecords"]
dataset = tf.data.TFRecordDataset(filenames)
def _parse_function(example_proto):
    keys_to_features = {'num_value':tf.VarLenFeature(tf.int64),
                        'list_value':tf.VarLenFeature(tf.int64)}
    parsed_features = tf.parse_single_example(example_proto, keys_to_features)
    return tf.sparse.to_dense(parsed_features['num_value']), \
           tf.sparse.to_dense(parsed_features['list_value'])
# Parse the record into tensors.
dataset = dataset.map(_parse_function)
# Shuffle the dataset
dataset = dataset.shuffle(buffer_size=1)
# Repeat the input indefinitly
dataset = dataset.repeat()  
# Generate batches
dataset = dataset.padded_batch(3, padded_shapes=([None],[None]), padding_values=(tf.constant(-1, dtype=tf.int64)
                                                 ,tf.constant(-1, dtype=tf.int64)))
iterator = dataset.make_one_shot_iterator()
i, data = iterator.get_next()

#Remove padding
data2= tf.RaggedTensor.from_tensor(data, padding=-1)

#Print data
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print(sess.run([ data, data2 ]))
    print(sess.run([ data, data2 ]))
    print(sess.run([ data, data2 ]))

Here is the official Tensorflow Guide to ragged tensors

https://www.tensorflow.org/guide/ragged_tensors

And official Tensorflow documentation

https://www.tensorflow.org/versions/r1.13/api_docs/python/tf/RaggedTensor

Solution

  • As you discovered, the RaggedTensors are not actually getting flattened. Internally, a 2D RaggedTensor is encoded using two Tensors/arrays: one containing a flat list of values, and the other containing row splits. For more details about how RaggedTensors are encoded using underlying tensors/arrays, see: https://www.tensorflow.org/guide/ragged_tensors#raggedtensor_encoding

    The confusion was probably coming from the way that RaggedTensors get displayed when printing. Python has two string conversion methods: __str__ and __repr__. __str__ gets used if you just print a value by itself, and __repr__ gets used if that value is embedded in some larger structure (such as a list).

    For RaggedTensorValue, the __str__ method returns "<tf.RaggedTensorValue %s>" % self.to_list(). I.e., it will show you the value formatted as a list. But the __repr__ method returns "tf.RaggedTensorValue(values=%r, row_splits=%r)" % (self._values, self._row_splits). I.e., it will show you the underlying numpy arrays that are used to encode the RaggedTensorValue.