Part 21: PyTorch, Debugging Autograd

Debugging issues with autograd can sometimes be challenging, especially when dealing with complex models and computations. Understanding common pitfalls and using the right debugging tools can help identify and resolve issues effectively.

Common Pitfalls

In-Place Operations:

• Problem: In-place operations can inadvertently modify tensors required for gradient computation, leading to errors.
• Solution: Avoid in-place operations on tensors that require gradients.

import torch

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
# In-place operation (Bad practice)
x.add_(1)  # Avoid this

# Non in-place operation (Good practice)
x = x + 1

Missing requires_grad=True:

• Problem: Forgetting to set requires_grad=True on tensors that need gradients.
• Solution: Ensure that all tensors involved in gradient computation have requires_grad=True.

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)  # Ensure this is set

No torch.no_grad() for Inference:

• Problem: Not using torch.no_grad() during inference leads to unnecessary memory usage.
• Solution: Use torch.no_grad() to disable gradient tracking during inference.

with torch.no_grad():
    outputs = model(inputs)
    print("Inference outputs:", outputs)

Using torch.autograd for Debugging

PyTorch provides several tools within the torch.autograd module to help debug and understand the autograd process.

• Using torch.autograd.gradcheck and torch.autograd.gradgradcheck: These functions check if the gradients are correct for custom Function implementations.
  • Gradient Check Example:

from torch.autograd import gradcheck

class MyReLU(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return input.clamp(min=0)

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        grad_input[input < 0] = 0
        return grad_input

input = torch.randn(5, requires_grad=True, dtype=torch.double)
relu = MyReLU.apply

test = gradcheck(relu, (input,), eps=1e-6, atol=1e-4)
print("Gradient check passed:", test)

Visualizing the Computation Graph

Visualizing the computation graph can help understand the flow of operations and identify issues.

• Visualization Tools:
  • PyTorch does not provide a built-in graph visualization tool, but external tools like TensorBoard or torchviz can be used.
  • Using torchviz for Visualization:

from torchviz import make_dot

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x + 2
z = y * y * 3
out = z.mean()

make_dot(out, params={"x": x})

Example: Debugging a Simple Model

Here’s an example of debugging a simple model using gradient checks and avoiding common pitfalls.

• Complete Debugging Example:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import gradcheck

class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(10, 50)
        self.fc2 = nn.Linear(50, 1)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# Initialize model, loss function, and optimizer
model = SimpleNN()
loss_fn = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# Create synthetic data
inputs = torch.randn(64, 10, requires_grad=True)
targets = torch.randn(64, 1)

# Training loop with gradient check
for epoch in range(100):
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = loss_fn(outputs, targets)
    
    if epoch % 10 == 0:
        # Perform gradient check every 10 epochs
        input_check = inputs.clone().detach().requires_grad_(True)
        output_check = model(input_check)
        test = gradcheck(lambda x: model(x), (input_check,), eps=1e-6, atol=1e-4)
        print(f"Epoch [{epoch}/100], Loss: {loss.item():.4f}, Gradient check passed: {test}")
    
    loss.backward()
    optimizer.step()

print("Training and debugging completed.")