GPU Acceleration in PyTorch

Hello! Let's continue our journey into PyTorch.

In our last lesson, you learned how to implement the five essential steps of a single training iteration: the forward pass, loss computation, gradient zeroing, backward pass, and weight update. You now understand the fundamental "engine" that makes a neural network learn from a batch of data. However, all our computations so far have been on the CPU. For the large models and datasets used in deep learning, this is often too slow.

This lesson focuses on a critical practical skill: accelerating your training by using a Graphics Processing Unit (GPU). Our learning outcome is to be able to move tensors and models to a GPU device in PyTorch. By the end of this lesson, you will know how to write "device-agnostic" code that seamlessly uses a GPU when available, drastically speeding up your model training.

Why Use a GPU?

At a high level, a CPU is composed of a few, very powerful cores optimized for sequential, complex tasks. A GPU, on the other hand, is composed of thousands of simpler cores designed for parallel computations. Training a neural network involves many independent matrix multiplications and vector operations, which is exactly the kind of parallel workload GPUs excel at.

PyTorch leverages NVIDIA's CUDA platform to perform computations on GPUs. Let's learn how to make our code use this power.

To begin, let's get a solid conceptual understanding of how PyTorch handles CPUs and GPUs.

The Core Mechanics: Device and .to()

As you've just read, the key to managing where your computations happen in PyTorch is the concept of a device. There are two main steps to using a GPU:

1. Define your target device: Create a torch.device object. It's best practice to do this dynamically, so your code can run on machines with or without a GPU.
2. Move your objects to the device: Use the .to() method to move your model and your data tensors to the selected device.

The following video from deeplizard provides a detailed walkthrough of these mechanics, including common errors you might encounter.

To summarize the key points from the reading and video:

• Device Object: The best practice is to define your device at the start of your script:

  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  print(f"Using device: {device}")
  

• Moving Objects: The .to(device) method is the primary tool. It returns a new copy of the object on the target device. The original object is not modified.
  • model = model.to(device)
  • input_tensor = input_tensor.to(device)
• Device Consistency: All tensors involved in an operation, including the model's parameters and the input data, must be on the same device. Failure to ensure this is one of the most common errors for beginners.

Modifying the Training Step for GPU Acceleration

Now, let's apply this to the training step we built in the last lesson. Modifying a training loop for GPU execution requires just three key changes.

The article "PyTorch in One Hour" provides a very clear summary of exactly which lines to change.

Let's integrate these three changes into the code from our previous lesson.

import torch
import torch.nn as nn




# --- Setup from previous lessons ---

class SimpleMLP(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, num_classes)
        )

    def forward(self, x):
        return self.network(x)

input_dim = 784
hidden_dim = 512
class_count = 10
learning_rate = 0.001




# --- GPU MODIFICATION 1: Define the device ---
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")





# Instantiate model, loss, and optimizer
model = SimpleMLP(input_size=input_dim, hidden_size=hidden_dim, num_classes=class_count)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)




# --- GPU MODIFICATION 2: Move the model to the device ---
# This is done once, after the model is created.
model.to(device)





# Create dummy data for one batch
batch_size = 4
X_batch = torch.randn(batch_size, input_dim)
y_batch = torch.randint(0, class_count, (batch_size,))

print(f"\nModel is on: {next(model.parameters()).device}")
print(f"Data is on: {X_batch.device}")




# --- GPU MODIFICATION 3: Move the data to the device ---
# This must be done for each batch inside the training loop.
X_batch = X_batch.to(device)
y_batch = y_batch.to(device)

print(f"Data has been moved to: {X_batch.device}")

print("\n--- Starting One Training Step on", device, "---")




# ----- THE TRAINING STEP SEQUENCE (now on the GPU) -----




# Step 1: Forward Pass (model and data are on the same device)
y_pred_logits = model(X_batch)




# Step 2: Compute Loss
loss = loss_fn(y_pred_logits, y_batch)
print(f"Calculated Loss: {loss.item():.4f}")




# Step 3: Zero Gradients
optimizer.zero_grad()




# Step 4: Backward Pass
loss.backward()




# Step 5: Update Weights
optimizer.step()

print("--- Finished One Training Step ---")

Important Practical Considerations

1. Data Transfer Overhead: Moving data from CPU memory to GPU memory (.to('cuda')) takes time. If your model is very small and you process tiny batches of data, this transfer time can actually make your GPU-based training slower than CPU training. The speedup from GPUs becomes significant for larger models and larger batch sizes where the computation time dwarfs the data transfer time.

2. Interacting with other libraries (like NumPy): NumPy arrays live on the CPU. You cannot directly convert a GPU tensor to a NumPy array. You must first move the tensor back to the CPU.

   
   
   
   # Assume 'gpu_tensor' is a tensor on the GPU
   # This will cause an error:
   # numpy_array = gpu_tensor.numpy() 
   
   
   
   
   # This is the correct way:
   cpu_tensor = gpu_tensor.cpu()
   numpy_array = cpu_tensor.numpy()
   
   
   
   
   # Or in one line:
   numpy_array = gpu_tensor.cpu().numpy()
   

   This is a common requirement when you want to save your results or use libraries like Matplotlib or Scikit-learn for plotting and evaluation.

Conclusion

You have now learned one of the most important practical skills for any deep learning practitioner: how to leverage a GPU to accelerate training.

Key Takeaways:

• GPUs accelerate deep learning by performing the massive number of parallel operations required for training large models.
• The standard practice for writing portable PyTorch code is to define a device object that automatically selects the GPU if available: device = torch.device("cuda" if torch.cuda.is_available() else "cpu").
• A model's parameters and the input data tensors must all be on the same device to perform computations.
• The process involves three changes to your training script:
  1. Define the device object.
  2. Move your model to the device once, after creating it (model.to(device)).
  3. Move each batch of data to the device inside your training loop (X_batch.to(device), y_batch.to(device)).

In this lesson and the previous one, we've been working with a single, manually created batch of data. To train a model properly, we need to efficiently load, shuffle, and iterate over an entire dataset. In our next lesson, we will learn how to use PyTorch's Dataset and DataLoader classes to manage this process professionally.