Building NN in PyTorch
PyTorch provides the `nn` module to build neural networks with layers, loss functions, and optimizers. This chapter shows how to create a simple multi‑layer perceptron (MLP) for MNIST classification.
Define a Model (Subclass `nn.Module`)
import torch.nn as nn
import torch.nn.functional as F
class MLP(nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(28*28, 128)
self.fc2 = nn.Linear(128, 64)
self.fc3 = nn.Linear(64, 10)
def forward(self, x):
x = x.view(-1, 28*28) # flatten
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return xTraining Loop
model = MLP()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
for epoch in range(5):
for batch_idx, (data, target) in enumerate(train_loader):
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()Evaluation
correct = 0
total = 0
with torch.no_grad():
for data, target in test_loader:
output = model(data)
_, predicted = torch.max(output, 1)
total += target.size(0)
correct += (predicted == target).sum().item()
print(f"Accuracy: {100*correct/total:.2f}%")Two Minute Drill
- Subclass `nn.Module` and implement `forward`.
- Use `nn.Linear` for dense layers, `F.relu` for activation.
- Loss function: `nn.CrossEntropyLoss` for classification.
- Optimizer: `torch.optim.Adam` or SGD.
- Training loop: zero gradients, forward, loss, backward, step.
Need more clarification?
Drop us an email at career@quipoinfotech.com