Python
TensorFlow
IO operations
machine learning
data processing

Python_io in tensorflow

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

tf.io is the TensorFlow module for reading files, decoding raw bytes, and serializing examples into formats that work efficiently with tf.data. It matters because input pipelines are often the real bottleneck in training jobs, and tf.io gives you TensorFlow-native building blocks instead of relying on ad hoc Python file handling inside the training loop.

Read and Write Files with TensorFlow Ops

The simplest tf.io functions work directly with file contents. They are useful when you want the pipeline to stay inside TensorFlow instead of dropping back to plain Python code.

python
1import tensorflow as tf
2
3path = "/tmp/example.txt"
4tf.io.write_file(path, "hello from tensorflow\n")
5contents = tf.io.read_file(path)
6
7print(contents.numpy().decode("utf-8"))

That example is small, but the main benefit is composability. The result of tf.io.read_file is a tensor, so it can flow straight into TensorFlow decoding ops and Dataset.map.

Decode Structured Bytes

tf.io is also where many format decoders live. Image pipelines are a common example.

python
1import tensorflow as tf
2
3image_bytes = tf.io.read_file("/tmp/sample.png")
4image = tf.io.decode_png(image_bytes, channels=3)
5image = tf.image.resize(image, [128, 128])
6
7print(image.shape)

The key idea is that reading bytes and decoding those bytes are separate steps. That separation lets you swap decoders for PNG, JPEG, WAV, or serialized examples while keeping the same dataset structure.

Serialize Data as TFRecord

TFRecord is TensorFlow's standard binary record format for large training datasets. tf.io provides the feature-building helpers used to encode and parse those records.

python
1import tensorflow as tf
2
3def _bytes_feature(value: bytes) -> tf.train.Feature:
4    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
5
6def _int64_feature(value: int) -> tf.train.Feature:
7    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
8
9example = tf.train.Example(features=tf.train.Features(feature={
10    "text": _bytes_feature(b"hello"),
11    "label": _int64_feature(1),
12}))
13
14serialized = example.SerializeToString()
15
16with tf.io.TFRecordWriter("/tmp/data.tfrecord") as writer:
17    writer.write(serialized)

Once records are written, you parse them with tf.io.parse_single_example.

python
1import tensorflow as tf
2
3feature_spec = {
4    "text": tf.io.FixedLenFeature([], tf.string),
5    "label": tf.io.FixedLenFeature([], tf.int64),
6}
7
8dataset = tf.data.TFRecordDataset(["/tmp/data.tfrecord"])
9dataset = dataset.map(lambda x: tf.io.parse_single_example(x, feature_spec))
10
11for item in dataset.take(1):
12    print(item["text"].numpy(), item["label"].numpy())

Why tf.io Beats Plain Python in Pipelines

Reading files with open() is fine for quick scripts, but it does not integrate as cleanly with graph execution, parallel dataset mapping, and TensorFlow's input pipeline optimizations. tf.io operations can be composed inside tf.data, parallelized with num_parallel_calls, and moved closer to the training runtime without rewriting the whole pipeline later.

That does not mean plain Python should never be used. It means data-loading logic that sits on the hot path of training usually benefits from TensorFlow-native operations.

Another practical benefit is deployment symmetry. If preprocessing depends on tf.io and other TensorFlow ops, the same logic is easier to reuse in training jobs, exported preprocessing layers, and serving code paths. That reduces the drift that often appears when Python-only file parsing lives outside the model pipeline.

For larger datasets, pair tf.io work with Dataset.cache, prefetch, and parallel mapping. Good I/O primitives help, but the surrounding dataset configuration still determines whether the accelerator waits on input.

Common Pitfalls

  • Mixing heavy Python-side file I/O into Dataset.map and then wondering why the pipeline is slow.
  • Confusing file reading with decoding and trying to use raw bytes as if they were already parsed tensors.
  • Writing TFRecords without a matching parse spec for reading them back.
  • Ignoring shape and dtype after decode operations.
  • Overusing tf.py_function when a native tf.io operation already exists.

Summary

  • 'tf.io provides TensorFlow-native file, decoding, and serialization operations.'
  • 'tf.io.read_file and tf.io.write_file are the basic building blocks for file content tensors.'
  • Decode ops such as tf.io.decode_png turn raw bytes into usable tensors.
  • TFRecord support in tf.io helps build scalable training datasets.
  • Native TensorFlow I/O integrates better with tf.data than ad hoc Python file handling.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.