Practice Now
TensorFlow
Multi-GPU Training
Machine Learning
Model Training
Python Scripting

How does one train multiple models in a single script in TensorFlow when there are GPUs present?

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
Deep LearningTensorFlow & KerasMachine LearningPython

Introduction

To train multiple models in a single TensorFlow script with GPUs, you have several options: assign each model to a specific GPU using tf.device, limit GPU memory growth so multiple models can share a GPU, or train models sequentially by clearing the session between runs. The key challenge is GPU memory — a single model often claims all available VRAM. You must either partition GPUs across models, enable memory growth, or use virtual devices to split a physical GPU into logical ones.

Enable Memory Growth (Essential First Step)

python
1import tensorflow as tf
2
3# Without this, TensorFlow allocates ALL GPU memory for the first model
4# and the second model gets an OOM error
5gpus = tf.config.list_physical_devices('GPU')
6for gpu in gpus:
7    tf.config.experimental.set_memory_growth(gpu, True)
8
9print(f"Available GPUs: {len(gpus)}")

By default, TensorFlow pre-allocates the entire GPU memory. With set_memory_growth(True), it allocates only as much memory as needed and grows on demand, leaving room for other models.

Strategy 1: One Model Per GPU

python
1import tensorflow as tf
2
3gpus = tf.config.list_physical_devices('GPU')
4for gpu in gpus:
5    tf.config.experimental.set_memory_growth(gpu, True)
6
7# Model A on GPU 0
8with tf.device('/GPU:0'):
9    model_a = tf.keras.Sequential([
10        tf.keras.layers.Dense(256, activation='relu', input_shape=(784,)),
11        tf.keras.layers.Dense(10, activation='softmax')
12    ])
13    model_a.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
14
15# Model B on GPU 1
16with tf.device('/GPU:1'):
17    model_b = tf.keras.Sequential([
18        tf.keras.layers.Dense(512, activation='relu', input_shape=(784,)),
19        tf.keras.layers.Dense(128, activation='relu'),
20        tf.keras.layers.Dense(10, activation='softmax')
21    ])
22    model_b.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
23
24# Train both (sequentially — each on its assigned GPU)
25model_a.fit(x_train, y_train, epochs=5, batch_size=64)
26model_b.fit(x_train, y_train, epochs=5, batch_size=64)

Strategy 2: Virtual GPU Devices (Split One GPU)

python
1import tensorflow as tf
2
3gpus = tf.config.list_physical_devices('GPU')
4if gpus:
5    # Split one physical GPU into two virtual GPUs with 4GB each
6    tf.config.set_logical_device_configuration(
7        gpus[0],
8        [
9            tf.config.LogicalDeviceConfiguration(memory_limit=4096),
10            tf.config.LogicalDeviceConfiguration(memory_limit=4096),
11        ]
12    )
13
14logical_gpus = tf.config.list_logical_devices('GPU')
15print(f"Physical GPUs: {len(gpus)}, Logical GPUs: {len(logical_gpus)}")
16
17# Assign models to virtual GPUs
18with tf.device(logical_gpus[0].name):
19    model_a = tf.keras.Sequential([...])
20    model_a.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
21
22with tf.device(logical_gpus[1].name):
23    model_b = tf.keras.Sequential([...])
24    model_b.compile(optimizer='adam', loss='sparse_categorical_crossentropy')

Strategy 3: Sequential Training with Session Clearing

python
1import tensorflow as tf
2import gc
3
4configs = [
5    {'units': 128, 'lr': 0.001, 'name': 'small_model'},
6    {'units': 512, 'lr': 0.0005, 'name': 'large_model'},
7    {'units': 256, 'lr': 0.001, 'name': 'medium_model'},
8]
9
10results = {}
11
12for config in configs:
13    # Clear previous model from GPU memory
14    tf.keras.backend.clear_session()
15    gc.collect()
16
17    model = tf.keras.Sequential([
18        tf.keras.layers.Dense(config['units'], activation='relu', input_shape=(784,)),
19        tf.keras.layers.Dense(10, activation='softmax')
20    ])
21    model.compile(
22        optimizer=tf.keras.optimizers.Adam(learning_rate=config['lr']),
23        loss='sparse_categorical_crossentropy',
24        metrics=['accuracy']
25    )
26
27    history = model.fit(x_train, y_train, epochs=10, validation_split=0.2, verbose=1)
28    results[config['name']] = max(history.history['val_accuracy'])
29    model.save(f"models/{config['name']}.keras")
30
31print("Results:", results)

tf.keras.backend.clear_session() destroys the current TF graph and releases GPU memory, allowing the next model to use it.

Strategy 4: Parallel Training with Multiprocessing

python
1import multiprocessing as mp
2import tensorflow as tf
3
4def train_model(gpu_id, config, return_dict):
5    import os
6    os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id)
7
8    import tensorflow as tf
9    gpus = tf.config.list_physical_devices('GPU')
10    for gpu in gpus:
11        tf.config.experimental.set_memory_growth(gpu, True)
12
13    model = tf.keras.Sequential([
14        tf.keras.layers.Dense(config['units'], activation='relu', input_shape=(784,)),
15        tf.keras.layers.Dense(10, activation='softmax')
16    ])
17    model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
18
19    # Load data inside the subprocess
20    (x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
21    x_train = x_train.reshape(-1, 784).astype('float32') / 255.0
22
23    history = model.fit(x_train, y_train, epochs=5, verbose=0)
24    return_dict[config['name']] = history.history['accuracy'][-1]
25
26if __name__ == '__main__':
27    manager = mp.Manager()
28    return_dict = manager.dict()
29
30    configs = [
31        {'units': 128, 'name': 'model_a'},
32        {'units': 256, 'name': 'model_b'},
33    ]
34
35    processes = []
36    for i, config in enumerate(configs):
37        p = mp.Process(target=train_model, args=(i, config, return_dict))
38        processes.append(p)
39        p.start()
40
41    for p in processes:
42        p.join()
43
44    print("Results:", dict(return_dict))

Each subprocess gets its own GPU via CUDA_VISIBLE_DEVICES. This avoids memory conflicts entirely because each process has an isolated TensorFlow runtime.

Common Pitfalls

  • Not enabling memory growth before creating models: TensorFlow allocates all GPU memory by default. If you create model A without set_memory_growth(True), it claims 100% of VRAM and model B gets an out-of-memory error. Call set_memory_growth before any model or tensor creation.
  • Setting CUDA_VISIBLE_DEVICES after importing TensorFlow: TensorFlow reads CUDA_VISIBLE_DEVICES at import time. Setting it after import tensorflow has no effect. In multiprocessing, set the environment variable before importing TF in each subprocess.
  • Forgetting clear_session() between sequential models: Without clearing the session, previous model graphs and optimizer state remain in GPU memory. After training 3-4 models sequentially, the GPU runs out of memory. Call tf.keras.backend.clear_session() and gc.collect() between each model.
  • Using threads instead of processes for parallel GPU training: Python threads share the same TensorFlow runtime and GPU context. Two threads training simultaneously on the same GPU cause memory conflicts and nondeterministic errors. Use multiprocessing with separate CUDA_VISIBLE_DEVICES for true parallel training.
  • Exceeding virtual device memory limits: When splitting a GPU into virtual devices with LogicalDeviceConfiguration(memory_limit=4096), each model must fit within that limit. A model that needs 6GB fails even though the physical GPU has 12GB. Size the virtual partitions based on each model's actual memory requirements.

Summary

  • Enable set_memory_growth(True) on all GPUs before creating any models
  • Use tf.device('/GPU:N') to pin each model to a specific GPU when multiple GPUs are available
  • Use virtual GPU devices to split a single GPU between multiple models with fixed memory budgets
  • Clear the session with tf.keras.backend.clear_session() between sequential model training runs
  • Use multiprocessing with CUDA_VISIBLE_DEVICES for fully parallel training on separate GPUs
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
Basis Lab
No course covers this one yet

Basis Lab writes one for you from a sentence about what you want to be able to do, then teaches it and checks you understood. Your first course is free.

Build my 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
Codemia
© 2026 Codemia
Resources
Legal
Social
All Rights Reserved.