tensorflow py_func is handy but makes my training step very slow.

tensorflow

py_func

training performance

model optimization

deep learning

tensorflow py_func is handy but makes my training step very slow.

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Start Practicing Learn More

Introduction

tf.py_func (TF1) and tf.py_function (TF2) let you run arbitrary Python code inside a TensorFlow computation graph. They are convenient for custom preprocessing, metrics, or operations that lack TensorFlow equivalents. However, they force execution back to Python on the CPU, bypassing TF's graph optimizer, XLA compilation, and GPU acceleration. A single py_func in your training loop can drop throughput by 5-50x because it serializes execution through the Python GIL.

Why py_func Is Slow

python

1import tensorflow as tf
2import numpy as np
3
4# This runs in Python, not on GPU
5def my_custom_op(x):
6    return np.sqrt(x.numpy()) * 2  # Copies data to CPU, runs in Python
7
8# In the training pipeline
9dataset = dataset.map(lambda x: tf.py_function(my_custom_op, [x], tf.float32))

Each call to py_function:

Copies tensor data from GPU to CPU (device-to-host transfer)
Acquires the Python GIL (blocks all other Python threads)
Runs the Python function (not optimized, no vectorization)
Copies the result back to GPU (host-to-device transfer)
Breaks graph optimization (XLA cannot see through py_func)

The GPU sits idle waiting for the Python function to finish.

Measuring the Impact

python

1import time
2
3# Without py_func
4@tf.function
5def train_step_fast(x, y):
6    with tf.GradientTape() as tape:
7        pred = model(x, training=True)
8        loss = loss_fn(y, pred)
9    grads = tape.gradient(loss, model.trainable_variables)
10    optimizer.apply_gradients(zip(grads, model.trainable_variables))
11    return loss
12
13# With py_func in preprocessing
14def slow_preprocess(x):
15    return np.clip(x.numpy(), 0, 1).astype(np.float32)
16
17# Benchmark
18start = time.time()
19for batch in dataset.take(100):
20    train_step_fast(batch[0], batch[1])
21print(f"Without py_func: {time.time() - start:.2f}s")  # about 2s
22
23start = time.time()
24for batch in dataset_with_pyfunc.take(100):
25    train_step_fast(batch[0], batch[1])
26print(f"With py_func: {time.time() - start:.2f}s")  # about 15s

Fix 1: Replace with TensorFlow Ops

Most Python/NumPy operations have TensorFlow equivalents:

python

1# SLOW: py_function with numpy
2def preprocess_slow(image):
3    img = image.numpy()
4    img = np.clip(img / 255.0, 0, 1)
5    img = np.rot90(img, k=1)
6    return img.astype(np.float32)
7
8dataset = dataset.map(
9    lambda x: tf.py_function(preprocess_slow, [x], tf.float32)
10)
11
12# FAST: Pure TensorFlow ops (runs on GPU, graph-optimized)
13def preprocess_fast(image):
14    image = tf.cast(image, tf.float32) / 255.0
15    image = tf.clip_by_value(image, 0.0, 1.0)
16    image = tf.image.rot90(image, k=1)
17    return image
18
19dataset = dataset.map(preprocess_fast)

Common replacements:

NumPy / Python	TensorFlow Equivalent
`np.clip`	`tf.clip_by_value`
`np.where`	`tf.where`
`np.concatenate`	`tf.concat`
`np.reshape`	`tf.reshape`
`np.argmax`	`tf.argmax`
`np.random.shuffle`	`tf.random.shuffle`
`cv2.resize`	`tf.image.resize`
`scipy.ndimage.rotate`	`tf.image.rot90` / `tfa.image.rotate`

Fix 2: Move py_func to Data Pipeline

If you must use py_func, isolate it in the data pipeline with prefetching:

python

1# Move Python ops to the data pipeline, not the training loop
2dataset = (tf.data.Dataset.from_tensor_slices((x_train, y_train))
3    .map(lambda x, y: (
4        tf.py_function(custom_augment, [x], tf.float32),
5        y
6    ), num_parallel_calls=tf.data.AUTOTUNE)
7    .batch(32)
8    .prefetch(tf.data.AUTOTUNE)  # Overlap CPU preprocessing with GPU training
9)

prefetch(AUTOTUNE) runs the Python preprocessing on CPU while the GPU trains on the previous batch. num_parallel_calls=AUTOTUNE runs multiple py_function calls in parallel (limited by GIL for CPU-bound work, but helps for I/O-bound operations).

Fix 3: Use tf.numpy_function (TF2)

tf.numpy_function is the TF2 replacement that gives cleaner semantics:

python

1def augment(image):
2    # Receives a numpy array directly (no .numpy() needed)
3    image = np.flipud(image)
4    return image.astype(np.float32)
5
6# Set output shape explicitly
7result = tf.numpy_function(augment, [image], tf.float32)
8result.set_shape(image.shape)  # Required because shape is unknown after numpy_function

Still slow for the same reasons, but avoids the .numpy() conversion step.

Fix 4: Custom TensorFlow Op

For operations called millions of times, write a custom C++/CUDA op:

python

1# Register a custom op (C++ implementation)
2# my_custom_op.cc compiled as shared library
3
4import tensorflow as tf
5custom_module = tf.load_op_library('./my_custom_op.so')
6result = custom_module.my_fast_op(input_tensor)

This runs on GPU with full graph optimization. Only worth the effort for ops called in every training step on every batch.

Fix 5: Use TensorFlow Addons or tf.image

Many custom operations already exist in TensorFlow or its ecosystem:

python

1import tensorflow_addons as tfa
2
3# Instead of py_func for image rotation
4rotated = tfa.image.rotate(image, angles=0.5)
5
6# Instead of py_func for cutout augmentation
7augmented = tfa.image.random_cutout(image, mask_size=(16, 16))
8
9# tf.image has many built-in operations
10augmented = tf.image.random_flip_left_right(image)
11augmented = tf.image.random_brightness(augmented, max_delta=0.2)
12augmented = tf.image.random_contrast(augmented, lower=0.8, upper=1.2)

Fix 6: Process Offline

If the custom operation does not depend on training state, preprocess the entire dataset once and save:

python

1# Preprocess once, save to TFRecord
2import tqdm
3
4writer = tf.io.TFRecordWriter("preprocessed.tfrecord")
5for image, label in tqdm.tqdm(raw_dataset):
6    processed = expensive_python_preprocessing(image.numpy())
7    example = serialize_example(processed, label.numpy())
8    writer.write(example)
9writer.close()
10
11# Load preprocessed data. No py_func needed during training
12dataset = tf.data.TFRecordDataset("preprocessed.tfrecord")
13dataset = dataset.map(parse_example).batch(32).prefetch(tf.data.AUTOTUNE)

Common Pitfalls

py_func inside @tf.function: tf.py_function works inside @tf.function but disables graph optimization for that portion. The entire traced function may fall back to eager mode.
Unknown output shapes: After py_function, TensorFlow does not know the output shape. Call result.set_shape(...) explicitly, or downstream operations that need static shapes will fail.
GIL bottleneck: num_parallel_calls helps with I/O-bound Python functions but not CPU-bound ones. Python's GIL serializes CPU-intensive numpy operations across threads.
Memory copies: Each py_function call copies data between TensorFlow's memory and Python's memory. For large tensors (images, embeddings), this copy alone can be the bottleneck.
Using py_func for simple math: Operations like clipping, normalization, and type casting all have direct TF equivalents. Check tf.math, tf.image, and tf.signal before using py_function.

Summary

tf.py_func / tf.py_function forces execution back to Python CPU, bypassing GPU and graph optimization
Replace with native TensorFlow ops (tf.clip_by_value, tf.image.*, tf.math.*) whenever possible
If unavoidable, isolate py_function in the data pipeline with prefetch(AUTOTUNE) to overlap with GPU training
Set output shapes explicitly with set_shape() after py_function
For dataset-wide preprocessing, process offline and save to TFRecord
Check TensorFlow Addons (tfa) for advanced operations before writing custom Python functions