Tensorflow DataSet Sequence Example

Tensorflow DataSet Sequence Example

Writer

// An highlighted block
import tensorflow as tf

sequences = [[1, 2, 3], [4, 5, 1], [1, 2]]
label_sequences = [[0, 1, 0], [1, 0, 0], [1, 1]]
 
rec_file="sequence_ex_record"
def make_example(sequence, labels):
    # The object we return
    ex = tf.train.SequenceExample()
    # A non-sequential feature of our example
    sequence_length = len(sequence)
    ex.context.feature["length"].int64_list.value.append(sequence_length+1)
    # Feature lists for the two sequential features of our example
    fl_tokens = ex.feature_lists.feature_list["tokens"]
    fl_labels = ex.feature_lists.feature_list["labels"]
    for token, label in zip(sequence, labels):
        fl_tokens.feature.add().int64_list.value.append(token)
        fl_labels.feature.add().int64_list.value.append(label)
    return ex
 
#Write all examples into a TFRecords file
with open(rec_file,"w") as fp:
    writer = tf.python_io.TFRecordWriter(fp.name)
    for sequence, label_sequence in zip(sequences, label_sequences):
        ex = make_example(sequence, label_sequence)
        writer.write(ex.SerializeToString())
    writer.close()

Reader

// An highlighted block
import tensorflow as tf

sequences = [[1, 2, 3], [4, 5, 1], [1, 2]]
label_sequences = [[0, 1, 0], [1, 0, 0], [1, 1]]
 
rec_file="sequence_ex_record"

def _parse(ex):
    # Define how to parse the example
    context_features = {
        "length": tf.FixedLenFeature([], dtype=tf.int64)
    }
    sequence_features = {
        "tokens": tf.FixedLenSequenceFeature([], dtype=tf.int64),
        "labels": tf.FixedLenSequenceFeature([], dtype=tf.int64)
    }
     
    # Parse the example
    context_parsed, sequence_parsed = tf.parse_single_sequence_example(
        serialized=ex,
        context_features=context_features,
        sequence_features=sequence_features
    )
    return {"length": context_parsed["length"]}, {"tokens": sequence_parsed["tokens"],"labels": sequence_parsed["labels"] }
    
dataset = tf.data.TFRecordDataset([rec_file]).map(_parse)
iterator = dataset.make_one_shot_iterator()

# This is an op that gets the next element from the iterator
next_element = iterator.get_next()

with tf.Session() as sess:
    print sess.run(next_element)