Complete PyTorch Interview Guide
Questions
Explain how PyTorch's autograd works and why it's essential for training neural networks.
The Scenario
You are an ML engineer at a self-driving car company. You are training a new model to detect pedestrians in a video feed. However, the model’s weights are not updating. The loss is stuck at a high value, and the model is not learning anything.
You have already checked the usual suspects: the learning rate is not too high or too low, the data is correctly pre-processed, and the model architecture seems to be correct.
You suspect that there might be an issue with the gradient computation.
The Challenge
Explain how you would use PyTorch’s autograd engine to debug this problem. What are the key concepts of autograd that you would use, and how would you use them to identify the source of the problem?
A junior engineer might just say that `autograd` is magic. They might not be able to explain how the computation graph works or how to use the `requires_grad` and `.grad` attributes to debug a model.
A senior engineer would know that `autograd` is a powerful tool for debugging models. They would be able to explain how the computation graph is built and how to use it to trace the flow of gradients. They would also have a clear plan for how to use `autograd` to identify the source of the problem.
Step 1: Understand the Key Concepts of autograd
| Concept | Description |
|---|---|
requires_grad | A boolean attribute on a tensor that tells autograd whether to track operations on it. |
grad_fn | A function that is associated with a tensor that was created by an operation. It knows how to compute the gradients of that operation. |
.grad | An attribute on a tensor that stores the gradients of that tensor after .backward() has been called. |
Step 2: Verify the Computation Graph
The first step is to verify that the computation graph is being built correctly. We can do this by inspecting the grad_fn of the tensors in our model.
import torch
import torch.nn as nn
# ... (define your model and loss function) ...
# Perform a forward pass
inputs = torch.randn(1, 3, 224, 224)
labels = torch.randn(1, 10)
outputs = model(inputs)
loss = loss_fn(outputs, labels)
# Check the grad_fn of the loss
print(loss.grad_fn)If the grad_fn of the loss is None, it means that the computation graph is not being built correctly. This could be because you have detached a tensor from the graph somewhere, or because you are using a non-differentiable operation.
Step 3: Check the Gradients
The next step is to check the gradients of the model’s parameters. We can do this by calling .backward() on the loss and then inspecting the .grad attribute of the parameters.
# Compute the gradients
loss.backward()
# Check the gradients of the first layer
print(model.fc1.weight.grad)If the gradients are None or all zeros, it means that the gradients are not being propagated back to the early layers of the model. This could be due to the vanishing gradient problem, or it could be because you have detached a tensor from the graph.
Step 4: Use Hooks to Inspect Intermediate Gradients
If you are still not sure what is going on, you can use hooks to inspect the gradients of the intermediate layers of the model.
def my_hook(grad):
print(grad)
# Register a hook on the output of the first layer
h = model.fc1.register_hook(my_hook)
# ... (perform a forward and backward pass) ...
# Remove the hook
h.remove()By using these techniques, you can systematically debug any issue with the gradient computation in your model.
Practice Question
You are debugging a model and you find that the gradients of the early layers are all zeros. What is the most likely cause of this?