# How to Handle Out-of-Memory Errors When Training CNN Models on Large Images

> Learn to handle out-of-memory errors when training CNN models on large images. Reduce GPU memory usage with resizing, smaller batches, mixed precision, and cache clearing.

- Repository: [Microsoft/AI-For-Beginners](https://github.com/microsoft/AI-For-Beginners)
- Tags: how-to-guide
- Published: 2026-08-09

---

**Reduce GPU memory consumption when training CNNs on high-resolution images by resizing inputs, lowering batch sizes, enabling mixed-precision training with `torch.cuda.amp`, and clearing the cache with `torch.cuda.empty_cache()` between epochs.**

Training convolutional neural networks on large images often triggers out-of-memory (OOM) errors because each forward pass stores activation tensors for every layer while the backward pass retains gradients for backpropagation. The microsoft/AI-For-Beginners curriculum provides practical mitigation strategies in its computer vision lessons, specifically within the ConvNets notebooks. Learning to handle out-of-memory errors when training CNN models on large images is essential for running experiments on consumer-grade GPUs with limited VRAM.

## Why CNNs Exhaust GPU Memory on Large Images

Convolutional layers allocate feature maps proportional to input spatial dimensions, meaning a 1024×1024 image consumes significantly more memory than a 224×224 image during training. Each training step stores intermediate activations from the forward pass to compute gradients during backpropagation, quickly exhausting GPU memory when processing high-resolution datasets. The PyTorch implementation in `lessons/4-ComputerVision/07-ConvNets/ConvNetsPyTorch.ipynb` demonstrates this behavior in its standard training loop.

## Resize Images to Reduce Activation Memory

Down-sampling images before they reach the model is the most effective way to handle out-of-memory errors when training CNN models on large images. Smaller spatial dimensions produce proportionally smaller activation tensors, dramatically cutting memory usage without modifying the network architecture.

Use `torchvision.transforms.Resize` to standardize input dimensions on-the-fly:

```python
import torchvision

transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize((224, 224)),   # Reduce from original resolution

    torchvision.transforms.ToTensor(),
])

```

This transformation ensures only resized tensors occupy GPU memory during the forward pass, preventing the allocation of massive feature maps that correspond to the original image size.

## Reduce Batch Size and Optimize Data Loading

Lowering the batch size directly reduces the number of simultaneous samples processed in each iteration, cutting the memory required for intermediate tensors. Additionally, configuring the `DataLoader` with `pin_memory=True` enables asynchronous CPU-to-GPU transfers, preventing unnecessary memory duplication.

```python
batch_size = 16   # Reduce from default 64 or 32

train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=batch_size,
    shuffle=True,
    pin_memory=True,          # Speeds up CPU→GPU transfer

    num_workers=4,
)

```

According to the source code in `lessons/4-ComputerVision/07-ConvNets/ConvNetsPyTorch.ipynb`, this configuration prevents the naïve approach of loading all images into RAM before training, instead keeping only the current batch in GPU memory at any given time.

## Enable Mixed-Precision Training

Full-precision (`float32`) tensors occupy double the memory of half-precision (`float16`) values. Enabling automatic mixed precision with `torch.cuda.amp` halves memory per tensor and often accelerates computation on modern NVIDIA GPUs.

Wrap the forward pass in an autocast context and use `GradScaler` to maintain gradient precision:

```python
scaler = torch.cuda.amp.GradScaler()

for epoch in range(num_epochs):
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        
        with torch.cuda.amp.autocast():
            outputs = model(images)
            loss = criterion(outputs, labels)
        
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

```

This approach is particularly effective for the CNN implementations found in the microsoft/AI-For-Beginners curriculum.

## Clear Cached Memory Between Epochs

PyTorch retains cached memory after tensor deallocation to speed up future allocations, which can trigger OOM errors when the allocator cannot reclaim sufficient contiguous blocks. Explicitly clearing this cache after validation or between epochs frees memory for the next training phase.

```python
for epoch in range(num_epochs):
    # Training loop...

    
    # Validation loop...

    
    # Release cached memory after each epoch

    torch.cuda.empty_cache()

```

Inserting `torch.cuda.empty_cache()` at strategic points in your training script prevents the accumulation of unused cached tensors that would otherwise cause out-of-memory failures during subsequent iterations.

## Use Gradient Checkpointing for Deep Networks

Very deep CNNs generate numerous intermediate activations that strain memory limits. Gradient checkpointing trades computational overhead for memory savings by recomputing activations during the backward pass instead of storing them.

Apply `torch.utils.checkpoint` to specific blocks of layers:

```python
from torch.utils.checkpoint import checkpoint

def forward_with_checkpoint(x):
    # Wrap memory-intensive blocks

    return checkpoint(model.block, x)

```

This technique is recommended in [`lessons/4-ComputerVision/08-TransferLearning/TrainingTricks.md`](https://github.com/microsoft/AI-For-Beginners/blob/main/lessons/4-ComputerVision/08-TransferLearning/TrainingTricks.md) for training deeper architectures on hardware with limited VRAM.

## Manage Multi-GPU Environments

Running multiple models simultaneously on a single GPU causes accidental memory oversubscription. When multiple GPUs are available, explicitly set the device to isolate each training run:

```python
device = torch.device('cuda:0')   # Select specific GPU

model = model.to(device)
torch.cuda.set_device(device)

```

The TensorFlow counterpart in `lessons/4-ComputerVision/07-ConvNets/ConvNetsTF.ipynb` implements similar device management strategies for distributed training scenarios.

## Summary

- **Resize images** to 224×224 or 128×128 using `torchvision.transforms.Resize` to reduce activation tensor sizes before they reach the network.
- **Lower batch sizes** and configure `DataLoader` with `pin_memory=True` to minimize simultaneous GPU memory allocation.
- **Enable mixed-precision training** with `torch.cuda.amp.autocast()` and `GradScaler()` to halve memory consumption compared to float32.
- **Clear cached memory** explicitly using `torch.cuda.empty_cache()` after validation loops or epochs to reclaim fragmented GPU memory.
- **Apply gradient checkpointing** via `torch.utils.checkpoint` for deep networks, trading compute for memory by recomputing activations during backpropagation.
- **Isolate GPU devices** using `torch.cuda.set_device()` when running multiple experiments to prevent memory contention.

## Frequently Asked Questions

### How do I know if my batch size is too large for my GPU?

If you encounter a `RuntimeError: CUDA out of memory` immediately after the forward pass begins, your batch size is likely too large for the available VRAM. Start with a batch size of 16 or 8 for high-resolution images (512×512 and above), then gradually increase while monitoring `nvidia-smi` or `torch.cuda.memory_allocated()` until you reach 80-90% GPU utilization.

### Does mixed-precision training affect model accuracy?

Mixed-precision training typically maintains model accuracy comparable to full-precision training because `torch.cuda.amp.GradScaler()` dynamically scales loss values to preserve small gradient values during backpropagation. According to the microsoft/AI-For-Beginners source code, this approach is safe for most CNN architectures while providing significant memory savings.

### Why does my GPU still show out-of-memory errors after I reduce the batch size?

Residual cached memory from previous training runs or memory fragmentation can cause OOM errors even with smaller batches. Call `torch.cuda.empty_cache()` after each epoch and ensure no other processes occupy the GPU by checking `nvidia-smi`. Additionally, verify that your input images are actually being resized before reaching the model, as loading full-resolution images into the DataLoader consumes substantial CPU RAM that can indirectly affect GPU memory management.

### Can I use these techniques with TensorFlow instead of PyTorch?

Yes, the fundamental strategies apply to both frameworks. The microsoft/AI-For-Beginners repository provides `lessons/4-ComputerVision/07-ConvNets/ConvNetsTF.ipynb` which demonstrates TensorFlow equivalents such as mixed-precision training via `tf.keras.mixed_precision.Policy` and memory optimization through `tf.data` pipeline configurations. The specific function names differ, but the concepts of resizing inputs, reducing batch sizes, and clearing memory caches remain identical.