TensorFlow
Keras
CPU usage
performance optimization
machine learning configuration

How can I reduce the number of CPUs used by Tensorlfow/Keras?

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

TensorFlow will happily use multiple CPU threads unless you tell it otherwise. That is good for throughput, but bad when you are sharing a machine, running multiple experiments at once, or trying to keep resource usage predictable in CI or on a laptop. The practical fix is to set TensorFlow thread limits early, before the runtime fully initializes.

Understand What You Are Limiting

TensorFlow CPU usage is usually governed by thread pools, not by a simple "use exactly N cores" switch. The two most important controls are:

  • inter-op parallelism: how many separate ops can run in parallel
  • intra-op parallelism: how many threads one op can use internally

If you reduce both values, TensorFlow becomes less aggressive about using the whole machine.

Set Thread Limits in Python Before Heavy Work Starts

The most direct approach is to configure TensorFlow's threading API right after importing TensorFlow and before building or running major workloads.

python
1import tensorflow as tf
2
3tf.config.threading.set_inter_op_parallelism_threads(1)
4tf.config.threading.set_intra_op_parallelism_threads(2)
5
6model = tf.keras.Sequential([
7    tf.keras.layers.Input(shape=(10,)),
8    tf.keras.layers.Dense(16, activation="relu"),
9    tf.keras.layers.Dense(1)
10])
11
12model.compile(optimizer="adam", loss="mse")

This tells TensorFlow to keep one op-level worker and at most two threads inside an individual op.

Environment Variables Also Matter

Some deployments prefer environment variables because they can be set outside the Python code.

bash
1export OMP_NUM_THREADS=2
2export TF_NUM_INTEROP_THREADS=1
3export TF_NUM_INTRAOP_THREADS=2
4python train.py

This is useful when:

  • the training code is shared and you do not want to edit it
  • you are controlling resources in a container or scheduler
  • you want different limits per run

Set these before the Python process starts.

Keep the Order of Initialization Clean

Thread settings should happen early. If TensorFlow has already created thread pools by the time you set the limits, the results may not match your expectation.

A safe pattern is:

  1. set environment variables before launch, or
  2. import TensorFlow and immediately call the threading setters, then
  3. build datasets, models, and training loops

Do not bury the configuration deep inside training code after the runtime is already hot.

tf.data Pipelines Can Also Use CPU Heavily

Sometimes the model is not the main CPU consumer. The input pipeline is.

For example, aggressive parallel mapping can raise CPU usage even if TensorFlow op threads are limited.

python
1import tensorflow as tf
2
3
4def preprocess(x):
5    return x * 2
6
7
8dataset = tf.data.Dataset.range(1000)
9dataset = dataset.map(preprocess, num_parallel_calls=1)
10dataset = dataset.batch(32)

If you use AUTOTUNE everywhere, TensorFlow may intentionally use more CPU to maximize throughput. Lowering num_parallel_calls is often necessary when your real goal is reduced CPU load rather than peak performance.

Do Not Confuse CPU Limits With GPU Use

Even on GPU-enabled training, TensorFlow still uses CPUs for input pipelines, orchestration, and some ops. Lowering CPU thread counts does not disable GPU usage. It only reduces CPU-side parallelism.

That is helpful when a GPU job still overwhelms shared CPU resources.

A Small Reproducible Example

python
1import numpy as np
2import tensorflow as tf
3
4tf.config.threading.set_inter_op_parallelism_threads(1)
5tf.config.threading.set_intra_op_parallelism_threads(1)
6
7x = np.random.rand(256, 10).astype("float32")
8y = np.random.rand(256, 1).astype("float32")
9
10model = tf.keras.Sequential([
11    tf.keras.layers.Dense(32, activation="relu"),
12    tf.keras.layers.Dense(1)
13])
14
15model.compile(optimizer="adam", loss="mse")
16model.fit(x, y, epochs=2, batch_size=32, verbose=1)

This is a simple way to test that the configuration is being applied before using it in a larger job.

Tradeoffs You Should Expect

Lower CPU usage usually means:

  • slower training
  • slower preprocessing
  • more predictable machine behavior
  • less contention with other applications

That is often a good trade in shared environments. The mistake is expecting the same throughput after deliberately reducing parallelism.

Common Pitfalls

  • Setting thread limits after the TensorFlow runtime is already heavily initialized.
  • Reducing model threads but leaving tf.data parallelism wide open.
  • Expecting CPU limits to act like strict OS-level core pinning.
  • Forgetting that environment variables must be set before the Python process starts.
  • Interpreting reduced throughput as a bug rather than the expected result of lower parallelism.

Summary

  • Limit TensorFlow CPU use by setting inter-op and intra-op thread counts early.
  • Environment variables are useful when you want per-run control without changing code.
  • Watch tf.data pipeline parallelism, not just model execution threads.
  • Reduced CPU use usually trades away training speed for predictability.
  • The key is to configure TensorFlow before it initializes its main execution pools.

Related reading
Course
Intermediate
27 lessons
15 hours
DSA Fundamentals

Master algorithmic patterns and data structures through hands-on LeetCode-style problems - from arrays and hashing to dynamic programming and advanced graphs.

View the 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.