Python / Python Deep Learning and Neural Networks Interview Questions
How do you use GPUs in PyTorch and what are the key patterns for writing device-agnostic code?
PyTorch's device abstraction allows the same code to run on CPU, single GPU, or multiple GPUs with minimal changes. The fundamental operations are moving tensors to a device with .to(device) or .cuda(), and ensuring model and data tensors always reside on the same device before any computation.
A critical performance concept: CPU–GPU data transfers are expensive (PCIe bandwidth is limited vs. GPU memory bandwidth). Minimise them by loading data onto the GPU once per batch, pre-computing dataset statistics on CPU, and avoiding frequent tensor transfers inside the training loop.
import torch import torch.nn as nn # Device-agnostic code pattern device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f'Using: {device}') # cuda / mps / cpu # Move model to device model = nn.Linear(10, 5).to(device) # Move data to device in the training loop for X_batch, y_batch in loader: X_batch = X_batch.to(device, non_blocking=True) y_batch = y_batch.to(device, non_blocking=True) y_hat = model(X_batch) # ... # Check which device a tensor is on t = torch.randn(3) print(t.device) # cpu t_gpu = t.cuda() # or t.to('cuda:0') print(t_gpu.device) # cuda:0 # Apple Silicon device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu') # Memory diagnostics print(torch.cuda.memory_allocated() / 1e9, 'GB allocated') print(torch.cuda.max_memory_allocated() / 1e9, 'GB peak') torch.cuda.empty_cache() # release unused cached GPU memory # Multi-GPU: DistributedDataParallel (DDP) preferred over DataParallel model_ddp = nn.parallel.DistributedDataParallel(model, device_ids=[0, 1])
More Related questions...