M
ModelArchViz

Architecture

2 hidden layers · sigmoid activations · backprop

mlp.pyColab
1import torch
2import torch.nn as nn
3
4# Define a multilayer perceptron.
5class MLP(nn.Module):
6 def __init__(
7 self,
8 input_dim=784, # Number of input dimensions.
9 hidden_dim=128, # Number of hidden dimensions.
10 output_dim=10 # Number of output dimensions.
11 ):
12 super().__init__()
13
14 # Register affine projections; initialization does not transform tensors.
15 self.hidden1 = nn.Linear(input_dim, hidden_dim)
16 self.hidden2 = nn.Linear(hidden_dim, hidden_dim)
17 self.output = nn.Linear(hidden_dim, output_dim)
18
19 def forward(self, x):
20 # First hidden block: (batch, input_dim) -> (batch, hidden_dim).
21 h1_pre = self.hidden1(x) # (batch, input_dim) -> (batch, hidden_dim)
22 h1 = torch.sigmoid(h1_pre) # (batch, hidden_dim)
23
24 # Second hidden block keeps the hidden shape: (batch, hidden_dim).
25 h2_pre = self.hidden2(h1) # (batch, hidden_dim)
26 h2 = torch.sigmoid(h2_pre) # (batch, hidden_dim)
27
28 # Output block: (batch, hidden_dim) -> (batch, output_dim).
29 logits = self.output(h2) # (batch, hidden_dim) -> (batch, output_dim)
30 return logits
31
32# Train on a tiny synthetic classification batch.
33model = MLP(input_dim=4, hidden_dim=8, output_dim=2)
34train_inputs = torch.tensor(
35 [
36 [1.0, 0.0, 1.0, 0.0],
37 [0.0, 1.0, 0.0, 1.0],
38 ]
39) # -> (2, 4)
40train_targets = torch.tensor([0, 1]) # -> (2)
41criterion = nn.CrossEntropyLoss()
42optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
43
44# Fit the model for a few steps on the tiny dataset.
45for step in range(3):
46 optimizer.zero_grad()
47 logits = model(train_inputs) # (2, 4) -> (2, 2)
48 loss = criterion(logits, train_targets) # (2, 2), (2) -> scalar
49 loss.backward()
50 optimizer.step()
51
52# Keep the final scalar loss for inspection.
53final_loss = loss.item() # scalar