win-tfrecord.py

import numpy as np
import tensorflow as tf


SIZE = int(64e6)

def _bytes_feature(value):
  return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))


def decode(example):
  feature_keys = {'encoded': tf.FixedLenFeature((), tf.string)}
  features = tf.parse_single_example(example, features=feature_keys)
  return tf.decode_raw(features['encoded'], tf.int32)


def encode(array):
  return tf.train.Example(features=tf.train.Features(feature={
      'encoded': _bytes_feature([array.astype(np.int32).tobytes()]),
    })).SerializeToString()


with tf.python_io.TFRecordWriter('foo.tfrecord') as writer:
  writer.write(encode(np.ones([SIZE], np.int32)))


dataset = tf.data.TFRecordDataset('foo.tfrecord').map(decode)

data = dataset.make_one_shot_iterator().get_next()

with tf.Session() as session:
  print(session.run(data).shape)