TensorFlow GPU显存选项配置无效，出现OOM内存分配错误求助

Hey there, sorry to hear you're stuck with this GPU memory exhaustion issue—even after turning on allow_growth, it's still possible to hit OOM depending on your model size, batch configuration, and other factors. Let's walk through practical, prioritized solutions to get your training back on track:

1. Shrink Your Batch Size First #

This is the fastest fix for most OOM cases. The error mentions a tensor with shape [64,160,400], so your current batch size is 64. Try cutting it in half (to 32) or even quartering it (to 16) and see if that resolves the issue:

# Update your training loop or data loader's batch_size parameter
model.fit(train_data, batch_size=32, epochs=10, ...)

Smaller batches mean less data is loaded into GPU memory at once, which directly reduces your memory footprint.

2. Enable Mixed Precision Training #

TensorFlow's mixed precision lets you use 16-bit floats (float16) for most computations instead of 32-bit, cutting memory usage by ~50% with almost no loss in model accuracy. Here's how to set it up:

from tensorflow.keras.mixed_precision import set_global_policy

# Activate mixed precision
set_global_policy('mixed_float16')

This automatically converts compatible tensors to float16, and it can even speed up training on modern GPUs that support FP16 acceleration.

3. Clean Up Memory Fragmentation & Residuals #

Even with allow_growth, fragmented GPU memory can leave you with insufficient contiguous space for large tensors. Try explicitly clearing Keras sessions and re-enabling memory growth before initializing your model:

import tensorflow as tf

# Clear any leftover Keras state and reset GPU memory
tf.keras.backend.clear_session()
physical_gpus = tf.config.list_physical_devices('GPU')
if physical_gpus:
    tf.config.experimental.set_memory_growth(physical_gpus[0], True)

Run this code before defining or loading your model to ensure a clean memory state.

4. Simplify or Prune Your Model #

If your model is overly large (e.g., too many convolution filters, unnecessary dense layers), it'll naturally gobble up more memory. Try these tweaks:

• Reduce the number of filters in convolutional layers (e.g., from 64 to 32)
• Replace large dense layers with global average pooling (which uses far less memory)
• Use model pruning to remove redundant weights:

from tensorflow_model_optimization.sparsity import keras as sparsity

# Define a pruning schedule to gradually remove 50% of weights
pruning_schedule = sparsity.PolynomialDecay(
    initial_sparsity=0.0,
    final_sparsity=0.5,
    begin_step=1000,
    end_step=10000
)

# Apply pruning to your existing model
pruned_model = sparsity.prune_low_magnitude(your_base_model, pruning_schedule=pruning_schedule)

5. Limit GPU Memory Allocation #

If your GPU is shared with other processes, or you don't want TensorFlow to hog all available memory, you can cap the amount it uses:

physical_gpus = tf.config.list_physical_devices('GPU')
if physical_gpus:
    # Allocate only 70% of the GPU's total memory to TensorFlow
    tf.config.set_logical_device_configuration(
        physical_gpus[0],
        [tf.config.LogicalDeviceConfiguration(memory_limit=int(physical_gpus[0].memory_limit * 0.7))]
    )

6. Check for Memory Leaks #

Occasionally, repeated tensor creation or model reinitialization in your training loop can cause gradual memory leaks. To debug this:

• Ensure your model is initialized once, not inside a loop
• Use tf.debugging.experimental.enable_dump_debug_info to track tensor allocation and identify leaks

内容的提问来源于stack exchange，提问作者Jatin Ganhotra