{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# TFRecords 101\n\n#### Why to use TFRecords?\n\n`TPUs` have **eight cores** which act as eight independent workers. We can get data to each core more efficiently by splitting the dataset into multiple files or **shards**. This way, each core can grab an independent part of the data as it needs.\n\nThe most convenient kind of file to use for sharding in TensorFlow is a `TFRecord`. A TFRecord is a binary file that contains sequences of byte-strings. Data needs to be `serialized` (encoded as a byte-string) before being written into a TFRecord.\n\nThe most convenient way of serializing data in TensorFlow is to wrap the data with `tf.Example`. This is a record format based on Google's protobufs but designed for TensorFlow. It's more or less like a `dict` with some type annotations.\n\n<blockquote style=\"margin-right:auto; margin-left:auto; background-color: #ebf9ff; padding: 1em; margin:24px;\">\n    <strong>Walkthrough: Building a Dataset of TFRecords</strong><br>\nFor an end-to-end walkthrough of building a sharded dataset of TFRecords for labeled images, see <a href=\"https://www.kaggle.com/ryanholbrook/walkthrough-building-a-dataset-of-tfrecords\"><strong>this notebook</strong></a>!\n</blockquote>","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19"}},{"cell_type":"code","source":"import numpy as np\nimport matplotlib.pyplot as plt\n\nfrom sklearn.datasets import load_sample_image\n\nimport tensorflow as tf\nfrom tensorflow.data import Dataset, TFRecordDataset\nfrom tensorflow.data.experimental import TFRecordWriter\nfrom tensorflow.io import FixedLenFeature, VarLenFeature\nfrom tensorflow.train import BytesList, FloatList, Int64List, Example, Features, Feature\n\nfrom IPython.display import Image","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"PATH = '/kaggle/working/data.tfrecord'","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Serialization \n\nA TFRecord is a kind of file that TensorFlow uses to store binary data. TFRecords contain sequences of byte-strings. Here is a very simple TFRecord:","metadata":{}},{"cell_type":"code","source":"with tf.io.TFRecordWriter(path=PATH) as f:\n    f.write(b'123')    # write one record\n    f.write(b'xyz314') # write another record\n\nwith open(PATH, 'rb') as f:\n    print(f.read())","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"A TFRecord is a sequence of bytes, so we have to turn our data into byte-strings before it can go into a TFRecord. We can use `tf.io.serialize_tensor` to turn a tensor into a byte-string and `tf.io.parse_tensor` to turn it back. It's important to keep track of your tensor's datatype (in this case `tf.uint8`) since you have to specify it when parsing the string back to a tensor again.","metadata":{}},{"cell_type":"code","source":"x = tf.constant([[1, 2], [3, 4]], dtype=tf.uint8)\nprint('x:', x, '\\n')\n\nx_bytes = tf.io.serialize_tensor(x)\nprint('x_bytes:', x_bytes, '\\n')\n\nprint('x:', tf.io.parse_tensor(x_bytes, out_type=tf.uint8))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## tf.data\n\nSo how do we write a dataset as a TFRecord? If your dataset is composed of byte-strings, you can use `data.TFRecordWriter`. To read it back again, use `data.TFRecordsDataset`.","metadata":{}},{"cell_type":"code","source":"# Construct a small dataset\nds = Dataset.from_tensor_slices([b'abc', b'123'])\n\n# Write the dataset to a TFRecord\nwriter = TFRecordWriter(PATH)\nwriter.write(ds)\n    \n# Read the dataset from the TFRecord\nds_2 = TFRecordDataset(PATH)\nfor x in ds_2:\n    print(x)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"If your dataset is composed of tensors, serialize them first by mapping `tf.io.serialize_tensor` over the dataset. Then, when you read them back, map `tf.io.parse_tensor` to turn the byte-strings back into tensors.","metadata":{}},{"cell_type":"code","source":"# Create a dataset\nfeatures = tf.constant([\n    [1, 2],\n    [3, 4],\n    [5, 6],\n], dtype=tf.uint8)\nds = Dataset.from_tensor_slices(features)\n\n# Serialize the tensors\nds_bytes = ds.map(tf.io.serialize_tensor)\n\n# Write a TFRecord\nwriter = TFRecordWriter(PATH)\nwriter.write(ds_bytes)\n\n# Read it back\nds_bytes_2 = TFRecordDataset(PATH)\nds_2 = ds_2.map(lambda x: tf.io.parse_tensor(x, out_type=tf.uint8))\n\n# They are the same!\nfor x in ds:\n    print(x)\nprint()\nfor x in ds_2:\n    print(x)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Serializing Images\n\nImages can be encoded in several ways:\n- **raw** encode with `tf.io.serialize_tensor`, decode with `tf.io.parse_tensor`\n- **jpeg** encode with `tf.io.encode_jpeg`, decode with `tf.io.decode_jpeg` or `tf.io.decode_and_crop_jpeg`\n- **png** encode with `tf.io.encode_png`, decode with `tf.io.decode_png`\n\nJust be sure to use whichever decoder goes with the encoder you chose. Generally, using jpeg encoding for images is a good idea when using TPUs since this can compress the data some, potentially improving data transfer time.","metadata":{}},{"cell_type":"code","source":"# Load numpy array\nimage_raw = load_sample_image('flower.jpg')\nprint(\"Type {} with dtype {}\".format(type(image_raw), image_raw.dtype))\nplt.imshow(image_raw)\nplt.title(\"Numpy\")\nplt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# jpeg encode / decode\nimage_jpeg = tf.io.encode_jpeg(image_raw)\nprint(\"Type {} with dtype {}\".format(type(image_jpeg), image_jpeg.dtype))\nprint(\"Sample: {}\".format(image_jpeg.numpy()[:25]))\nImage(image_jpeg.numpy())","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"image_raw_2 = tf.io.decode_jpeg(image_jpeg)\n\nprint(\"Type {} with dtype {}\".format(type(image_raw_2), image_raw_2.dtype))\nplt.imshow(image_raw_2)\nplt.title(\"Numpy\")\nplt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## tf.Example\n\nWhat if you have structured data, like `(image, label)` pairs? TensorFlow also includes an API for structured data, `tf.Example`. They are based on Google's [Protocol Buffers](https://developers.google.com/protocol-buffers).\n\nA single `Example` is meant to represent a single instance in a dataset, like a single `(image, label)` pair. Each `Example` has `Features`, described as a `dict` of feature names and values. A value can be either a `BytesList`, a `FloatList`, or an `Int64List`, each wrapped as a single `Feature`. There's no value type for tensors; instead, serialize tensors with `tf.io.serialize_tensor`, get the bytestring with the `numpy` method, and encode them in a `BytesList`.\n\nHere's how we could encode labeled image data:","metadata":{}},{"cell_type":"code","source":"# The Data\nimage = tf.constant([ # this could also be a numpy array\n    [0, 1, 2],\n    [3, 4, 5],\n    [6, 7, 8],\n])\nlabel = 0\nclass_name = \"Class A\"\n\n\n# Wrap with Feature as a BytesList, FloatList, or Int64List\nimage_feature = Feature(\n    bytes_list=BytesList(value=[\n        tf.io.serialize_tensor(image).numpy(),\n    ])\n)\nlabel_feature = Feature(\n    int64_list=Int64List(value=[label]),\n)\nclass_name_feature = Feature(\n    bytes_list=BytesList(value=[\n        class_name.encode()\n    ])\n)\n\n\n# Create a Features dictionary\nfeatures = Features(feature={\n    'image': image_feature,\n    'label': label_feature,\n    'class_name': class_name_feature,\n})\n\n# Wrap with Example\nexample = Example(features=features)\n\nprint(example)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"All of the data is stored as attributes of the `Example` instance.","metadata":{}},{"cell_type":"code","source":"print(example.features.feature['label'])","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Once everything is encoded as an `Example`, you can serialize it with the `SerializeToString` method.","metadata":{}},{"cell_type":"code","source":"example_bytes = example.SerializeToString()\nprint(example_bytes)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"It's nice to wrap all this in a function.","metadata":{}},{"cell_type":"code","source":"def make_example(image, label, class_name):\n    image_feature = Feature(\n        bytes_list=BytesList(value=[\n            tf.io.serialize_tensor(image).numpy(),\n        ])\n    )\n    label_feature = Feature(\n        int64_list=Int64List(value=[\n            label,\n        ])\n    )\n    class_name_feature = Feature(\n        bytes_list=BytesList(value=[\n            class_name.encode(),\n        ])\n    )\n\n    features = Features(feature={\n        'image': image_feature,\n        'label': label_feature,\n        'class_name': class_name_feature,\n    })\n    \n    example = Example(features=features)\n    \n    return example.SerializeToString()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"example = make_example(\n    image=np.array([[1, 2], [3, 4]]),\n    label=1,\n    class_name=\"Class B\",\n)\n\nprint(example)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The whole process might look something like:\n1. Build a dataset with `tf.data.Dataset`. You could use the `from_generator` or `from_tensor_slices` methods.\n2. Serialize the dataset by iterating over the dataset with `make_example`.\n3. Write the dataset to TFRecords with `io.TFRecordWriter` or `data.TFRecordWriter`.\n\nNote, however, that to use a function like `make_example` with the `Dataset` map method you'll need to wrap it with `tf.py_function` first since TensorFlow executes dataset transformations in graph mode. You could write something like this:\n\n```python\nds_bytes = ds.map(lambda image, label: tf.py_function(func=make_example, inp=[image, label], Tout=tf.string))\n```\n\nSee the [API docs](https://www.tensorflow.org/api_docs/python/tf/data/Dataset#map) for more info.\n\nFor TPU training, you'll want to shard the dataset into multiple files. the `Dataset.shard` method is handy for this, and Kaggle's [TPU documentation](https://www.kaggle.com/docs/tpu#tpu3) gives some advice about constructing shards.","metadata":{}},{"cell_type":"markdown","source":"## Parsing Serialized Examples\n\nTo decode a serialized example, we need to give TensorFlow a description of what kind of data to expect. We have just scalar entries for each so we can use `FixedLenFeature`. Pass the description to `tf.io.parse_single_example` along with the serialized example.","metadata":{}},{"cell_type":"code","source":"feature_description = {\n    'image': FixedLenFeature([], tf.string),\n    'label': FixedLenFeature([], tf.int64),\n    'class_name': FixedLenFeature([], tf.string),\n}\n\nexample_2 = tf.io.parse_single_example(example, feature_description)\nprint(\"Parsed:   \", example_2)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"These are the functions from the helper script that assemble the dataset from the TFRecords. We haven't covered everything here, but hopefully it's a little more clear what's going on.","metadata":{}},{"cell_type":"code","source":"def decode_image(image_data):\n    image = tf.image.decode_jpeg(image_data, channels=3)\n    image = tf.cast(image, tf.float32) / 255.0  # convert image to floats in [0, 1] range\n    image = tf.reshape(image, [*IMAGE_SIZE, 3]) # explicit size needed for TPU\n    return image\n\n\ndef read_labeled_tfrecord(example):\n    LABELED_TFREC_FORMAT = {\n        \"image\": tf.io.FixedLenFeature([], tf.string), # tf.string means bytestring\n        \"class\": tf.io.FixedLenFeature([], tf.int64),  # shape [] means single element\n    }\n    example = tf.io.parse_single_example(example, LABELED_TFREC_FORMAT)\n    image = decode_image(example['image'])\n    label = tf.cast(example['class'], tf.int32)\n    return image, label # returns a dataset of (image, label) pairs\n\n\ndef load_dataset(filenames, labeled=True, ordered=False):\n    # Read from TFRecords. For optimal performance, reading from multiple files at once and\n    # disregarding data order. Order does not matter since we will be shuffling the data anyway.\n\n    ignore_order = tf.data.Options()\n    if not ordered:\n        ignore_order.experimental_deterministic = False # disable order, increase speed\n\n    dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTO) # automatically interleaves reads from multiple files\n    dataset = dataset.with_options(ignore_order) # uses data as soon as it streams in, rather than in its original order\n    dataset = dataset.map(read_labeled_tfrecord if labeled else read_unlabeled_tfrecord, num_parallel_calls=AUTO)\n    # returns a dataset of (image, label) pairs if labeled=True or (image, id) pairs if labeled=False\n    return dataset","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## More Resources\n\n- [TPU-speed data pipelines: tf.data.Dataset and TFRecords](https://codelabs.developers.google.com/codelabs/keras-flowers-data/#4)\n- [Protocol Buffers](https://developers.google.com/protocol-buffers/)\n- [TFRecord and tf.Example](https://www.tensorflow.org/tutorials/load_data/tfrecord)\n\n## Reference\n\nThis notebook is replica of this source [notebook](https://www.kaggle.com/ryanholbrook/tfrecords-basics)\n\n---","metadata":{}}]}