Deep learning with PyTorch: your first model
Build and train your first neural network in PyTorch on a real dataset: tensors, DataLoader, model, loss, optimiser, the training loop, evaluation, saving and using a GPU. Runs free in Colab or Kaggle.
What you will build: a neural network in PyTorch that predicts whether a tumour is malignant or benign from 30 measurements (scikit-learn's built-in breast cancer dataset). It is a small, real, classic problem that trains in seconds on a CPU.
Prerequisites: the lesson *Neural networks explained*, basic pandas/NumPy, and a free notebook in Google Colab or Kaggle (see *Free notebooks: Google Colab and Kaggle*). PyTorch and scikit-learn are already installed in both. On your own laptop, install with pip install torch scikit-learn (check pytorch.org for the command for your system).
This dataset is for learning only. Never use a model like this for real medical decisions.
Tensors: PyTorch's arrays
A tensor is like a NumPy array that can also run on a GPU and remember how it was calculated, so PyTorch can compute gradients for you.
import torch
a = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
print(a.shape, a.dtype)
print(a * 2)
print(a @ a) # matrix multiplication
w = torch.tensor(3.0, requires_grad=True)
loss = (w * 2 - 10) ** 2 # a tiny function of w
loss.backward() # compute gradients automatically
print("d loss / d w =", w.grad) # 2 * (2w - 10) * 2 = -16 when w = 3This is autograd: you write the forward calculation and PyTorch gives you the gradient.
torch.Size([2, 2]) torch.float32
tensor([[2., 4.],
[6., 8.]])
tensor([[ 7., 10.],
[15., 22.]])
d loss / d w = tensor(-16.)The plan
- 1Open a new Colab or Kaggle notebook and paste each code cell below in order, running it before moving on.
- 2Prepare the data (Step 1) and wrap it in a DataLoader (Step 2).
- 3Define the model, loss and optimiser (Step 3).
- 4Run the training loop (Step 4), then evaluate (Step 5).
- 5Save the weights and try a prediction (Step 6).
Step 1: prepare the data
Neural networks train badly when features have very different scales (one column in thousands, another below 1). So we standardise each feature, fitting the scaler on the training data only so no information leaks from the test set.
import torch
import torch.nn as nn
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
torch.manual_seed(42)
X, y = load_breast_cancer(return_X_y=True)
print(X.shape, y.shape, "positive rate:", y.mean().round(3))
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=42, stratify=y
)
scaler = StandardScaler().fit(X_train) # learn mean/std from train only
X_train = scaler.transform(X_train)
X_test = scaler.transform(X_test)
# Convert to float32 tensors. Labels get shape (n, 1) to match the model output.
X_train = torch.tensor(X_train, dtype=torch.float32)
X_test = torch.tensor(X_test, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.float32).unsqueeze(1)
y_test = torch.tensor(y_test, dtype=torch.float32).unsqueeze(1)569 rows, 30 features. In this dataset label 1 means benign and 0 means malignant.
Step 2: batches with a DataLoader
Instead of updating weights once per epoch using all rows, we feed small batches. This is faster per update and the noise helps training. TensorDataset pairs features with labels, and DataLoader shuffles and slices them.
from torch.utils.data import TensorDataset, DataLoader
train_ds = TensorDataset(X_train, y_train)
train_loader = DataLoader(train_ds, batch_size=32, shuffle=True)
xb, yb = next(iter(train_loader))
print(xb.shape, yb.shape)torch.Size([32, 30]) torch.Size([32, 1])Step 3: define the model
nn.Sequential chains layers. Our network is 30 inputs, one hidden layer of 16 neurons with ReLU, and 1 output. The output is a raw score called a logit. We do not add a sigmoid inside the model, because the loss function we use applies it for us in a numerically safer way.
model = nn.Sequential(
nn.Linear(30, 16),
nn.ReLU(),
nn.Linear(16, 1),
)
print(model)
loss_fn = nn.BCEWithLogitsLoss() # binary cross-entropy on logits
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)BCEWithLogitsLoss is the standard loss for yes/no problems. Adam is a popular optimiser: gradient descent with automatic step-size adjustment.
Step 4: the training loop
Every PyTorch training loop has the same five lines inside it. Learn them by heart:
1. optimizer.zero_grad(): clear old gradients (they otherwise accumulate).
2. logits = model(xb): forward pass.
3. loss = loss_fn(logits, yb): measure the error.
4. loss.backward(): backpropagation.
5. optimizer.step(): update the weights.
for epoch in range(30):
model.train()
running_loss = 0.0
for xb, yb in train_loader:
optimizer.zero_grad()
logits = model(xb)
loss = loss_fn(logits, yb)
loss.backward()
optimizer.step()
running_loss += loss.item() * xb.size(0)
avg_loss = running_loss / len(train_ds)
if (epoch + 1) % 5 == 0:
print(f"epoch {epoch + 1:2d} train loss {avg_loss:.4f}")The loss should fall steadily each time it is printed.
epoch 5 train loss 0.1...
epoch 10 train loss 0.08...
(the exact values vary a little; they should keep shrinking)Step 5: evaluate on data the model has not seen
In evaluation we switch off gradient tracking with torch.no_grad() (faster, less memory) and call model.eval() (matters for layers like dropout, harmless here). A sigmoid turns the logit into a probability, and we call anything above 0.5 class 1.
model.eval()
with torch.no_grad():
test_logits = model(X_test)
probs = torch.sigmoid(test_logits)
preds = (probs > 0.5).float()
accuracy = (preds == y_test).float().mean().item()
print(f"test accuracy: {accuracy:.3f}")Accuracy on this dataset is typically in the mid-to-high 90s percent. Exact values depend on the seed and PyTorch version.
Reuse scikit-learn metrics
Accuracy is not the only metric. For anything medical or fraud-related, check precision and recall too (see *Metrics: accuracy, precision, recall, F1 and ROC*). You can pass preds.numpy() and y_test.numpy() straight into scikit-learn's classification_report.
Step 6: save, load and predict
torch.save(model.state_dict(), "model.pt") # save only the weights
# Later (new session): rebuild the same architecture, then load the weights
new_model = nn.Sequential(nn.Linear(30, 16), nn.ReLU(), nn.Linear(16, 1))
new_model.load_state_dict(torch.load("model.pt"))
new_model.eval()
with torch.no_grad():
one = X_test[:1] # one patient, shape (1, 30)
p = torch.sigmoid(new_model(one)).item()
print("probability of class 1:", round(p, 3), " actual:", int(y_test[0].item()))state_dict holds only numbers, so you must recreate the architecture before loading. Remember to also save the scaler (with joblib) if you want to predict on raw data.
Using a GPU
On Colab choose Runtime, then Change runtime type, then a GPU option. Then move the model and each batch to the GPU. For a model this tiny a GPU brings no benefit, but the pattern matters for bigger models.
device = "cuda" if torch.cuda.is_available() else "cpu"
print("using", device)
model = model.to(device)
for xb, yb in train_loader:
xb, yb = xb.to(device), yb.to(device)
# ... same five training lines as beforeBoth the model and the data must be on the same device, or PyTorch raises a RuntimeError.
Common mistakes
Beginner bugs
| Mistake | Symptom | Fix |
|---|---|---|
| Forgetting zero_grad() | Loss behaves strangely because gradients pile up | Call optimizer.zero_grad() every batch |
| Applying sigmoid and then BCEWithLogitsLoss | Poor training, because the sigmoid is applied twice | Feed raw logits to BCEWithLogitsLoss |
| Labels with wrong shape or type | Shape or dtype error | Use float32 labels of shape (n, 1) for this loss |
| Not scaling features | Slow or unstable training | Use StandardScaler, fit on training data only |
| Evaluating without no_grad() | Wasted memory | Wrap evaluation in torch.no_grad() |
| Model and data on different devices | RuntimeError | Move both with .to(device) |
I am learning PyTorch. My training code gives this error:
[paste the full error message]
Here is my code:
[paste code]
Explain in simple words what the error means, which line causes it, and how to fix it. Tell me the shape and dtype each tensor should have at that line.What to expect: A good answer points to a specific line and tells you to print .shape and .dtype of your tensors. Run the fix yourself rather than trusting it blindly.
Try it: Improve the model
Change the network to two hidden layers (30 to 32 to 16 to 1), and train for 50 epochs. Then try adding nn.Dropout(0.2) after each ReLU. Record the test accuracy for each version. Next, add a validation split and print validation loss each epoch to see when it stops improving.
Key takeaways
- A PyTorch workflow is: tensors, Dataset and DataLoader, model, loss, optimiser, training loop, evaluation.
- The training loop is always zero_grad, forward, loss, backward, step.
- Scale features and fit the scaler on training data only.
- Use model.eval() and torch.no_grad() when predicting.
- Save weights with state_dict, and keep the architecture code and the scaler too.
- Move model and data to the same device when using a GPU.
Check yourself
Finishing the quiz marks the lesson as done.
1. What does loss.backward() do?
2. Why do we call optimizer.zero_grad() in each iteration?
3. You use nn.BCEWithLogitsLoss. What should the final layer of the model output?
4. Why is the StandardScaler fitted on the training set only?
5. What is the purpose of torch.no_grad() during evaluation?
Saved on this device. Create a free account to keep it on every device.
Practise this with Sam
A real mock interview, out loud, with an honest scorecard. Free to start.