How to Avoid "CUDA Out of Memory" in PyTorch

When working with PyTorch and large deep learning models, especially on GPU (CUDA), running into the dreaded "CUDA out of memory" error is common. This issue can disrupt training, inference, or testing, particularly when dealing with large datasets or complex models.

In this article, we’ll explore several techniques to help you avoid this error and ensure your training runs smoothly on the GPU.

Introduction to CUDA Out of Memory Error

The "CUDA out of memory" error occurs when your GPU does not have enough memory to allocate for the task. PyTorch attempts to allocate memory dynamically, but if the memory demand exceeds the available capacity, you’ll see an error like this:

import torch

# Check if CUDA is available
if torch.cuda.is_available():
    device = torch.device("cuda")
    print("CUDA is available! Using GPU.")

    try:
        # Allocate a large tensor on the GPU (this size will likely exceed your GPU memory)
        large_tensor = torch.randn((100000, 10000, 10000), device=device)
        print("Tensor created successfully!")
    except RuntimeError as e:
        # Catch the CUDA Out of Memory error
        print(f"Caught a RuntimeError: {e}")
else:
    print("CUDA is not available. Please run this code on a system with a GPU.")

Output:

CUDA is available! Using GPU.
Caught a RuntimeError: CUDA out of memory. Tried to allocate 37252.90 GiB. GPU 0 has a total capacity of 14.75 GiB of which 14.65 GiB is free. Process 5534 has 100.00 MiB memory in use. Of the allocated memory 0 bytes is allocated by PyTorch, and 0 bytes is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. 

Techniques to Avoid CUDA Out of Memory Error

1. Reduce Batch Size

One of the most straightforward solutions is to reduce the batch size. Batch size refers to the number of samples processed before the model updates its weights. A larger batch size consumes more memory, and reducing it can often free up significant memory.

Example: If you are currently using a batch size of 64, try reducing it to 32 or even 16. Experiment with different sizes to find a balance between memory usage and training speed.

2. Use torch.cuda.empty_cache()

PyTorch does not release GPU memory after each operation. Instead, it reuses the allocated memory for future operations. You can manually clear unused GPU memory with the torch.cuda.empty_cache() function. This command does not reset the allocated memory but frees the cache for other parts of your program.

Usage: Call this function at appropriate places in your code, such as at the end of each epoch.

import torch
torch.cuda.empty_cache()

3. Use Gradient Accumulation

Gradient accumulation allows you to split a large batch into smaller sub-batches. After computing the loss and gradients for each sub-batch, accumulate the gradients over several steps before performing a weight update. This approach reduces memory usage while keeping the effective batch size large.

Implementation: Instead of reducing batch size, accumulate gradients over n steps to achieve the same effect as training with a larger batch size.

optimizer.zero_grad()

for i in range(grad_accum_steps):
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    loss = loss / grad_accum_steps
    loss.backward()

optimizer.step()

4. Utilize Mixed Precision Training

Mixed precision training involves using both 16-bit and 32-bit floating-point numbers to perform computations. By doing so, it reduces memory consumption and can speed up training. PyTorch provides the torch.cuda.amp (Automatic Mixed Precision) package for easy implementation.

Usage: Wrap your training loop with PyTorch’s AMP utilities.

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for inputs, labels in data_loader:
    with autocast():
        outputs = model(inputs)
        loss = criterion(outputs, labels)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

5. Use Model Parallelism

If your GPU does not have enough memory, consider distributing the model across multiple GPUs. PyTorch allows you to split your model and load parts of it onto different devices.

Example: Splitting the model across two GPUs:

model.layer1.to('cuda:0')
model.layer2.to('cuda:1')

This approach ensures that each GPU handles only a portion of the model, thereby reducing the memory burden on each device.

6. Clear Unused Variables with del

In some cases, you may be holding on to tensors that are no longer needed. You can manually release memory by using Python’s del statement to delete variables you no longer need. Combined with torch.cuda.empty_cache(), this method can free up memory.

Example:

del some_tensor
torch.cuda.empty_cache()

7. Check for Memory Leaks

A common source of the "CUDA out of memory" error is a memory leak caused by creating new variables inside loops without freeing the old ones. Ensure that any variable that you no longer need is explicitly deleted or goes out of scope.

Advanced Memory Management Techniques

1. Profile Memory Usage with torch.cuda.memory_summary()

PyTorch provides built-in functions to profile GPU memory usage. Use torch.cuda.memory_summary() to track how much memory is being used at different points in your code. This can help identify inefficient memory usage patterns or leaks.

Usage:

print(torch.cuda.memory_summary())

2. Use torch.no_grad() for Inference

When performing inference or evaluation, gradients are not needed, yet PyTorch computes them by default. Wrapping your inference code in torch.no_grad() prevents unnecessary memory allocation for gradients, freeing up GPU memory.

Example:

with torch.no_grad():
    outputs = model(inputs)

Conclusion

The "CUDA out of memory" error is a common hurdle when training large models or handling large datasets. However, with strategies such as reducing batch size, using gradient accumulation, mixed precision training, and more, you can often prevent this issue and make better use of your GPU resources. It's all about balancing memory usage with model complexity and computational needs.

By applying these techniques, you can keep your PyTorch training on track and make the most of your available GPU memory.

deepkumarpatra

Follow

Improve

Article Tags :

• Deep Learning
• AI-ML-DS
• Python-PyTorch
• AI-ML-DS With Python