tfrecord

When using MultiTFRecordDataset or TFRecordDataset in PyTorch, you can control the following behaviors:

• Finiteness: By default, MultiTFRecordDataset is infinite. Set infinite=False to make it finite.
• Shuffling: Both datasets automatically shuffle data if you provide a shuffle_queue_size argument.

# Finite dataset
dataset = MultiTFRecordDataset(..., infinite=False)

# Shuffled dataset
dataset = TFRecordDataset(..., shuffle_queue_size=1024)

You can pass a transform function to TFRecordDataset to perform post-processing (like image decoding or normalization) on features before they are returned. This is useful for converting raw bytes into usable tensors or arrays.

import tfrecord
import cv2

def decode_image(features):
    # get BGR image from bytes
    features["image"] = cv2.imdecode(features["image"], -1)
    return features

description = {
    "image": "bytes",
}

dataset = tfrecord.torch.TFRecordDataset("/tmp/data.tfrecord",
                                         index_path=None,
                                         description=description,
                                         transform=decode_image)

data = next(iter(dataset))
print(data)

To use the library with PyTorch IterableDataset capabilities, install the package with the [torch] extra:

pip3 install 'tfrecord[torch]'

It is highly recommended to create an index file (*.tfindex) for each TFRecord file. Providing an index file is required when using multiple workers to prevent the loader from returning duplicate records.

To create an index file for a single TFRecord file:

python3 -m tfrecord.tools.tfrecord2idx <tfrecord path> <index path>

To create index files for all *.tfrecord files in a directory:

tfrecord2idx <data dir>

To read SequenceExample records, use tfrecord.tfrecord_loader and provide both context_description and sequence_description. The loader will yield a tuple of (context_features, sequence_features).

import tfrecord

context_description = {"length": "int", "label": "int"}
sequence_description = {"tokens": "int", "seq_labels": "int"}
loader = tfrecord.tfrecord_loader("/tmp/data.tfrecord", None,
                                  context_description,
                                  sequence_description=sequence_description)

for context, sequence_feats in loader:
    print(context["label"])
    print(sequence_feats["seq_labels"])

Use MultiTFRecordDataset to sample from multiple TFRecord files based on provided probabilities. This class accepts patterns for both the TFRecord files and their corresponding index files.

import torch
from tfrecord.torch.dataset import MultiTFRecordDataset

tfrecord_pattern = "/tmp/{}.tfrecord"
index_pattern = "/tmp/{}.index"
splits = {
    "dataset1": 0.8,
    "dataset2": 0.2,
}
description = {"image": "byte", "label": "int"}

dataset = MultiTFRecordDataset(tfrecord_pattern, index_pattern, splits, description)
loader = torch.utils.data.DataLoader(dataset, batch_size=32)

data = next(iter(loader))
print(data)

To write SequenceExample records, use tfrecord.TFRecordWriter.write with two dictionaries: the first for context features and the second for sequence features. Sequence features should be provided as lists/sequences.

import tfrecord

writer = tfrecord.TFRecordWriter("/tmp/data.tfrecord")
# First dict is context, second dict is sequence
writer.write({'length': (3, 'int'), 'label': (1, 'int')},
             {'tokens': ([[0, 0, 1], [0, 1, 0], [1, 0, 0]], 'int'), 'seq_labels': ([0, 1, 1], 'int')})
writer.close()

Use TFRecordDataset to create a PyTorch-compatible dataset from a single TFRecord file. You must provide a description dictionary mapping feature names to types (e.g., "byte", "float", "int").

import torch
from tfrecord.torch.dataset import TFRecordDataset

tfrecord_path = "/tmp/data.tfrecord"
index_path = None
description = {"image": "byte", "label": "float"}

dataset = TFRecordDataset(tfrecord_path, index_path, description)
loader = torch.utils.data.DataLoader(dataset, batch_size=32)

data = next(iter(loader))
print(data)

Use tfrecord.TFRecordWriter to create TFRecord files. When writing, provide a dictionary where values are tuples containing the data and its type (e.g., (value, "type")).

import tfrecord

writer = tfrecord.TFRecordWriter("/tmp/data.tfrecord")
writer.write({
    "image": (image_bytes, "byte"),
    "label": (label, "float"),
    "index": (index, "int")
})
writer.close()

Use tfrecord.tfrecord_loader to iterate over records in a TFRecord file. It requires the file path, an optional index path, and a description dictionary.

import tfrecord

loader = tfrecord.tfrecord_loader("/tmp/data.tfrecord", None, {
    "image": "byte",
    "label": "float",
    "index": "int"
})
for record in loader:
    print(record["label"])