Introduction

When storing training data in TFRecord format, integer arrays are usually written with an int64_list feature. This is common for token ids, class id sequences, and sparse categorical values. The main goal is creating examples that are easy to serialize, parse, and feed into a tf.data pipeline without type mismatches.

Build Integer Features for tf.train.Example

TensorFlow uses protocol buffer messages for TFRecord entries. For integer lists, wrap values in tf.train.Int64List.

1import tensorflow as tf
2
3
4def int_list_feature(values):
5    return tf.train.Feature(int64_list=tf.train.Int64List(value=list(values)))
6
7
8def create_example(token_ids, label_id):
9    features = tf.train.Features(
10        feature={
11            "token_ids": int_list_feature(token_ids),
12            "label_id": int_list_feature([label_id]),
13        }
14    )
15    return tf.train.Example(features=features)
16
17example = create_example([10, 42, 7, 99], 1)
18print(example)

Use Python int values or NumPy integer types that convert cleanly to int64.

Write TFRecord Files

Once examples are created, serialize each example and write bytes with TFRecordWriter.

1import tensorflow as tf
2
3records = [
4    create_example([1, 3, 5], 0),
5    create_example([2, 4, 6, 8], 1),
6    create_example([11, 12], 0),
7]
8
9output_path = "dataset.tfrecord"
10with tf.io.TFRecordWriter(output_path) as writer:
11    for ex in records:
12        writer.write(ex.SerializeToString())
13
14print("Wrote:", output_path)

This produces a binary file that is efficient for streaming during model training.

Parse Integer List Features During Input

Parsing is where many mistakes happen. FixedLenFeature is for fixed-length arrays, while VarLenFeature handles variable-length lists.

1import tensorflow as tf
2
3feature_spec = {
4    "token_ids": tf.io.VarLenFeature(tf.int64),
5    "label_id": tf.io.FixedLenFeature([1], tf.int64),
6}
7
8
9def parse_record(serialized):
10    parsed = tf.io.parse_single_example(serialized, feature_spec)
11    token_ids = tf.sparse.to_dense(parsed["token_ids"])
12    label = parsed["label_id"][0]
13    return token_ids, label
14
15
16dataset = tf.data.TFRecordDataset(["dataset.tfrecord"])\
17    .map(parse_record)\
18    .padded_batch(2, padded_shapes=([None], []))
19
20for batch_tokens, batch_labels in dataset.take(1):
21    print(batch_tokens)
22    print(batch_labels)

If list lengths vary, use padded_batch so tensors in a batch share shape.

Keep Schema Explicit and Stable

Treat feature names and types as a schema contract between writers and readers. If you rename a key or switch dtype without updating the parser, failures can be silent until training diverges.

A robust pattern is storing schema constants in one module used by both write and read code paths.

FEATURE_TOKEN_IDS = "token_ids"
FEATURE_LABEL_ID = "label_id"

This reduces typo-driven bugs in larger pipelines.

Validate Round-Trip Data Integrity

After building a writer and parser, add a small round-trip check to confirm that values are preserved exactly. This is especially useful before launching expensive training jobs.

1import tensorflow as tf
2
3raw = [
4    ([9, 8, 7], 1),
5    ([100, 200], 0),
6]
7
8with tf.io.TFRecordWriter("check.tfrecord") as w:
9    for tokens, label in raw:
10        w.write(create_example(tokens, label).SerializeToString())
11
12read_back = list(tf.data.TFRecordDataset(["check.tfrecord"]).map(parse_record).as_numpy_iterator())
13print(read_back)

If parsed values differ from input, fix schema mismatches immediately. Round-trip tests catch feature key typos, dtype mismatches, and parsing bugs before they affect model metrics.

Production Pipeline Tip

For larger datasets, shard TFRecord outputs into multiple files such as train-00000-of-00010 style naming. Sharding improves parallel reads in distributed training jobs and reduces single-file bottlenecks. Keep shard counts proportional to dataset size so each file remains reasonably large but not unwieldy for upload and caching.

Common Pitfalls

One common issue is using tf.int32 in parsing while writing int64 lists. TensorFlow will raise type errors or require explicit casts.

Another pitfall is passing nested lists accidentally, for example a two-dimensional list into Int64List. That feature expects a flat sequence of integers.

People also use FixedLenFeature for variable-length sequences. This fails on examples with shorter or longer arrays.

A final issue is writing and reading from different feature names across scripts. Keep names centralized and test round-trip serialization early.

Summary

• Use tf.train.Int64List for integer list features in TFRecord examples.
• Serialize examples with TFRecordWriter and parse with matching schema.
• Use VarLenFeature for variable-length integer sequences.
• Convert sparse parsed values to dense tensors when needed.
• Keep feature names and dtypes stable across data pipeline stages.