# How to Fix Out‑of‑Memory Errors When Training GANs or Large CNNs in PyTorch > Resolve out of memory errors in PyTorch for GANs and CNNs. Learn to reduce batch size, use mixed precision, gradient accumulation, and clear CUDA caches for efficient training. - Repository: [Microsoft/AI-For-Beginners](https://github.com/microsoft/AI-For-Beginners) - Tags: how-to-guide - Published: 2026-08-20 --- **Out‑of‑memory (OOM) errors in deep learning are caused by GPU memory fragmentation from large tensors, model parameters, feature maps, and optimizer states, and can be systematically resolved by reducing batch size, enabling mixed‑precision training, using gradient accumulation, and clearing CUDA caches between epochs.** Training generative adversarial networks (GANs) and large convolutional neural networks (CNNs) frequently exhaust GPU memory—particularly on consumer‑grade hardware with 8–16 GB VRAM. The Microsoft **AI‑For‑Beginners** repository provides foundational notebooks for these architectures, but the reference implementations prioritize clarity over memory efficiency. This guide explains the root causes of OOM failures and provides production‑ready fixes you can apply directly to the repository's training loops. --- ## Why GANs and Large CNNs Exhaust GPU Memory PyTorch allocates GPU memory for four primary components during training: - **Input tensors** – batched images, labels, and intermediate data - **Model parameters** – weights and biases for all layers - **Feature maps** – activations stored for backpropagation - **Optimizer state** – momentum buffers, second‑order statistics for adaptive optimizers like Adam GANs compound this problem by maintaining **two separate networks** (generator and discriminator), each with their own parameters, gradients, and optimizer states. In `lessons/4-ComputerVision/10-GANs/GANPyTorch.ipynb`, the reference DC‑GAN implementation trains both networks simultaneously, effectively doubling memory pressure compared to a single CNN classifier. --- ## Immediate Fixes: Reduce Batch Size and Input Resolution The fastest mitigation is reducing data dimensions. Locate the dataset initialization in the affected notebook and apply both changes: ```python import torchvision.transforms as transforms # Original: batch_size = 64, images at full resolution (e.g., 128×128) batch_size = 32 # Reduce by 50% transform = transforms.Compose([ transforms.Resize(64), # Downsample from original resolution transforms.CenterCrop(64), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ]) ``` Reducing batch size from 64 to 32 and resolution from 128×128 to 64×64 cuts activation memory by approximately **75%** (quadratic spatial reduction × linear batch reduction). --- ## Enable Mixed‑Precision Training with `torch.cuda.amp` The **most impactful optimization** for modern NVIDIA GPUs (Volta architecture and newer) is Automatic Mixed Precision (AMP). AMP uses **FP16** (half‑precision) for forward/backward computations where safe, and **FP32** for critical reductions, halving activation memory without sacrificing convergence. Replace the standard training loop in `GANPyTorch.ipynb` with this AMP‑enabled version: ```python from torch.cuda.amp import autocast, GradScaler scaler = GradScaler() for epoch in range(num_epochs): for i, (real_images, _) in enumerate(dataloader): real_images = real_images.to(device) batch_size = real_images.size(0) # Labels for real and fake images real_labels = torch.ones(batch_size, 1).to(device) fake_labels = torch.zeros(batch_size, 1).to(device) # === Train Discriminator === optimizer_D.zero_grad() with autocast(): # Mixed‑precision forward outputs = discriminator(real_images) d_loss_real = criterion(outputs, real_labels) noise = torch.randn(batch_size, latent_dim, 1, 1).to(device) fake_images = generator(noise) outputs = discriminator(fake_images.detach()) d_loss_fake = criterion(outputs, fake_labels) d_loss = d_loss_real + d_loss_fake scaler.scale(d_loss).backward() # Scaled backward scaler.step(optimizer_D) # === Train Generator === optimizer_G.zero_grad() with autocast(): outputs = discriminator(fake_images) g_loss = criterion(outputs, real_labels) scaler.scale(g_loss).backward() scaler.step(optimizer_G) scaler.update() # Adjust loss scaling for next iteration # Critical: clear memory fragmentation after each epoch torch.cuda.empty_cache() ``` **Key components:** - `autocast()` – automatically selects FP16 for matmuls and convolutions - `GradScaler()` – scales loss to prevent gradient underflow in FP16 - `scaler.update()` – adjusts scale factor based on whether infinities/NaNs occurred --- ## Simulate Large Batches with Gradient Accumulation When research requires large effective batch sizes (e.g., 256 or 512) but hardware cannot accommodate them, **gradient accumulation** splits the batch into micro‑batches and accumulates gradients before the optimizer step. ```python accumulation_steps = 4 # Effective batch size = 4 × micro_batch_size micro_batch_size = 32 # Fits in GPU memory for epoch in range(num_epochs): optimizer_D.zero_grad() optimizer_G.zero_grad() for i, (real_images, _) in enumerate(dataloader): real_images = real_images.to(device) # Scale loss by accumulation steps to maintain gradient magnitude with autocast(): # ... compute d_loss and g_loss as before ... scaled_d_loss = d_loss / accumulation_steps scaled_g_loss = g_loss / accumulation_steps scaler.scale(scaled_d_loss).backward(retain_graph=True) scaler.scale(scaled_g_loss).backward() # Only update weights every N steps if (i + 1) % accumulation_steps == 0: scaler.step(optimizer_D) scaler.step(optimizer_G) scaler.update() optimizer_D.zero_grad() optimizer_G.zero_grad() torch.cuda.empty_cache() ``` This pattern achieves the statistical benefits of large‑batch training while keeping peak memory usage constant. --- ## Apply Gradient Checkpointing for Deep Networks For architectures deeper than the reference DC‑GAN, **gradient checkpointing** trades computation for memory by recomputing forward activations during backward passes instead of storing them. PyTorch implements this in `torch.utils.checkpoint`. Insert checkpointing in the generator or discriminator definition: ```python import torch.utils.checkpoint as checkpoint class DeepGenerator(nn.Module): def __init__(self): super().__init__() self.block1 = self._make_block(...) self.block2 = self._make_block(...) # Additional blocks... def forward(self, x): # Recompute block1 during backward pass x = checkpoint.checkpoint(self.block1, x) x = self.block2(x) # Normal execution return x ``` **Trade‑off:** Checkpointing increases training time by approximately **20–30%** but reduces activation memory by **40–60%** for deep residual networks. --- ## Memory‑Safe Validation with `torch.no_grad()` Validation and inference must explicitly disable gradient computation. In `lessons/4-ComputerVision/07-ConvNets/ConvNetsPyTorch.ipynb`, wrap evaluation as follows: ```python def validate(model, dataloader, criterion, device): model.eval() total_loss = 0.0 correct = 0 with torch.no_grad(): # Disables autograd engine for images, labels in dataloader: images, labels = images.to(device), labels.to(device) with autocast(): # Also use AMP for validation speedup outputs = model(images) loss = criterion(outputs, labels) total_loss += loss.item() _, predicted = torch.max(outputs, 1) correct += (predicted == labels).sum().item() return total_loss / len(dataloader), correct / len(dataloader.dataset) ``` `torch.no_grad()` prevents allocation of computation graphs, typically reducing validation memory by **30–50%**. --- ## Complete OOM‑Safe GAN Training Script Combine all techniques into a drop‑in replacement for the original `GANPyTorch.ipynb` training loop: ```python import torch import torch.nn as nn from torch.cuda.amp import autocast, GradScaler def train_oom_safe_gan( generator, discriminator, dataloader, num_epochs, latent_dim, device, batch_size=32, accumulation_steps=1 ): criterion = nn.BCELoss() optimizer_G = torch.optim.Adam(generator.parameters(), lr=0.0002, betas=(0.5, 0.999)) optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999)) scaler = GradScaler(init_scale=2**15) # Conservative initial scale generator.to(device) discriminator.to(device) for epoch in range(num_epochs): for i, (real_images, _) in enumerate(dataloader): real_images = real_images.to(device) current_batch = real_images.size(0) real_labels = torch.ones(current_batch, 1).to(device) fake_labels = torch.zeros(current_batch, 1).to(device) # === Discriminator === if i % accumulation_steps == 0: optimizer_D.zero_grad(set_to_none=True) #更高效地清除梯度 with autocast(): d_real = discriminator(real_images) loss_d_real = criterion(d_real, real_labels) noise = torch.randn(current_batch, latent_dim, 1, 1, device=device) with torch.no_grad(): # 生成假图时不构建图 fake_images = generator(noise) d_fake = discriminator(fake_images.detach()) loss_d_fake = criterion(d_fake, fake_labels) loss_d = (loss_d_real + loss_d_fake) / accumulation_steps scaler.scale(loss_d).backward() if (i + 1) % accumulation_steps == 0: scaler.step(optimizer_D) # === Generator === if i % accumulation_steps == 0: optimizer_G.zero_grad(set_to_none=True) with autocast(): g_output = discriminator(fake_images) loss_g = criterion(g_output, real_labels) / accumulation_steps scaler.scale(loss_g).backward() if (i + 1) % accumulation_steps == 0: scaler.step(optimizer_G) scaler.update() torch.cuda.empty_cache() # Monitor memory usage if device.type == 'cuda': allocated = torch.cuda.memory_allocated(device) / 1024**3 reserved = torch.cuda.memory_reserved(device) / 1024**3 print(f"Epoch [{epoch+1}/{num_epochs}] - Memory: {allocated:.2f}GB allocated, {reserved:.2f}GB reserved") ``` --- ## Key Files in Microsoft AI‑For‑Beginners | File | Purpose | Direct Link | |------|---------|-------------| | `lessons/4-ComputerVision/10-GANs/GANPyTorch.ipynb` | Reference DC‑GAN implementation | [View](https://github.com/microsoft/AI-For-Beginners/blob/main/lessons/4-ComputerVision/10-GANs/GANPyTorch.ipynb) | | `lessons/4-ComputerVision/07-ConvNets/ConvNetsPyTorch.ipynb` | CNN classifier basics | [View](https://github.com/microsoft/AI-For-Beginners/blob/main/lessons/4-ComputerVision/07-ConvNets/ConvNetsPyTorch.ipynb) | | `lessons/4-ComputerVision/08-TransferLearning/TransferLearningPyTorch.ipynb` | Layer freezing for memory reduction | [View](https://github.com/microsoft/AI-For-Beginners/blob/main/lessons/4-ComputerVision/08-TransferLearning/TransferLearningPyTorch.ipynb) | | `lessons/4-ComputerVision/09-Autoencoders/AutoEncodersPyTorch.ipynb` | Autoencoder memory stress test | [View](https://github.com/microsoft/AI-For-Beginners/blob/main/lessons/4-ComputerVision/09-Autoencoders/AutoEncodersPyTorch.ipynb) | --- ## Summary - **OOM errors** stem from the combined memory footprint of inputs, parameters, feature maps, and optimizer states—particularly severe in GANs with dual networks - **Reduce batch size and input resolution** as immediate, zero‑cost mitigations - **Mixed‑precision training** (`torch.cuda.amp`) halves activation memory on modern GPUs with automatic loss scaling - **Gradient accumulation** simulates large batches without proportional memory increase - **Gradient checkpointing** trades 20–30% compute for 40–60% memory savings in deep networks - **`torch.no_grad()` and `torch.cuda.empty_cache()`** eliminate unnecessary allocations and fragmentation - The Microsoft AI‑For‑Beginners repository's GAN and CNN notebooks benefit from systematic application of these techniques --- ## Frequently Asked Questions ### How do I know if my GPU supports mixed‑precision training? Mixed‑precision requires NVIDIA GPUs with **Tensor Cores** (Volta architecture and newer: V100, T4, RTX 20/30/40 series, A100, H100). Verify support by running `torch.cuda.is_available()` and checking your GPU compute capability—the `torch.cuda.amp` module will automatically fall back to FP32 on unsupported hardware without errors. ### Why does gradient accumulation require loss scaling? Without dividing by `accumulation_steps`, gradients from accumulated micro‑batches sum to N× their intended magnitude, causing optimizer steps that are too large and training instability. The division ensures the effective gradient matches what would result from a single large batch. ### When should I use `torch.cuda.empty_cache()` versus restarting the kernel? `torch.cuda.empty_cache()` releases unused cached memory to the system allocator but does not defragment the GPU memory pool. Use it between epochs or after large temporary allocations. If `nvidia-smi` shows substantial "reserved" memory with low "allocated" memory despite `empty_cache()` calls, a kernel restart may be necessary to fully defragment. ### Does gradient checkpointing slow down training significantly? Checkpointing increases forward pass time by approximately **20–30%** because activations must be recomputed during backward passes. However, it enables training models that would otherwise OOM, making the trade‑off essential for deep architectures beyond typical consumer GPU limits.