PyTorch CNN vs MLP
A CNN vs MLP comparison trains both model kinds on the same data and scores them on the same tests. A shifted copy of the validation set then shows what each model really learned.
Count both models' parameters
The MLP pays one weight per pixel for its first layer. The CNN shares each kernel over the whole image.
print("CNN:", sum(p.numel() for p in ShapesCNN().parameters()))
print("MLP:", sum(p.numel() for p in ShapesMLP().parameters()))The shift, in code
The shift copies the image into a blank frame, one row down and two columns right. The cells left behind stay zero.
x = torch.zeros(1, 1, 4, 4)
x[0, 0, 1, 1] = 1.0
print(shift_images(x, 1, 2)[0, 0])Same input, same output shape
Both models take the same image batch. Both return one score row per image.
batch = torch.zeros(2, 1, 16, 16)
print(ShapesCNN()(batch).shape, ShapesMLP()(batch).shape)Example
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
def draw_shape(label, g):
rows = torch.arange(16).view(16, 1).float()
cols = torch.arange(16).view(1, 16).float()
if label == 0:
centre = 8 + int(torch.randint(-2, 3, (1,), generator=g))
mask = (rows >= centre - 1) & (rows <= centre + 1) & (cols >= 0)
elif label == 1:
cr = 7.5 + int(torch.randint(-1, 2, (1,), generator=g))
cc = 7.5 + int(torch.randint(-1, 2, (1,), generator=g))
dist = torch.sqrt((rows - cr) ** 2 + (cols - cc) ** 2)
mask = (dist >= 4.5) & (dist <= 6.5)
else:
top = int(torch.randint(2, 7, (1,), generator=g))
left = int(torch.randint(2, 7, (1,), generator=g))
mask = (rows >= top) & (rows <= top + 7) & (cols >= left) & (cols <= left + 7)
return mask.float()
def make_shapes(count, seed):
g = torch.Generator().manual_seed(seed)
labels = torch.arange(count) % 3
labels = labels[torch.randperm(count, generator=g)]
images = [draw_shape(int(label), g) for label in labels]
return torch.stack(images).unsqueeze(1), labels
def shift_images(images, dy, dx):
out = torch.zeros_like(images)
out[:, :, dy:, dx:] = images[:, :, :-dy, :-dx]
return out
class ShapesCNN(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 8, 3, padding=1)
self.pool = nn.MaxPool2d(2)
self.conv2 = nn.Conv2d(8, 16, 3, padding=1)
self.gap = nn.AdaptiveAvgPool2d(1)
self.head = nn.Linear(16, 3)
def forward(self, x):
x = self.pool(torch.relu(self.conv1(x)))
x = torch.relu(self.conv2(x))
x = self.gap(x)
return self.head(torch.flatten(x, 1))
class ShapesMLP(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Flatten(),
nn.Linear(256, 64),
nn.ReLU(),
nn.Linear(64, 3),
)
def forward(self, x):
return self.net(x)
train_x, train_y = make_shapes(240, seed=42)
val_x, val_y = make_shapes(60, seed=43)
loader = DataLoader(TensorDataset(train_x, train_y), batch_size=32,
shuffle=True, generator=torch.Generator().manual_seed(0))
loss_fn = nn.CrossEntropyLoss()
def train(model):
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=4, gamma=0.5)
for epoch in range(8):
model.train()
for xb, yb in loader:
optimizer.zero_grad()
loss = loss_fn(model(xb), yb)
loss.backward()
optimizer.step()
scheduler.step()
return model
def accuracy(model, x, y):
model.eval()
with torch.no_grad():
predictions = model(x).argmax(dim=1)
return float((predictions == y).float().mean())
torch.manual_seed(1)
cnn = train(ShapesCNN())
torch.manual_seed(1)
mlp = train(ShapesMLP())
shifted_x = shift_images(val_x, 3, 3)
print("Model Parameters Accuracy Shifted accuracy")
for name, model in [("CNN", cnn), ("MLP", mlp)]:
params = sum(p.numel() for p in model.parameters())
acc = accuracy(model, val_x, val_y)
shifted = accuracy(model, shifted_x, val_y)
print(f"{name:<7}{params:<12}{acc:.3f} {shifted:.3f}")Key points
- Parameter counts: 1,299 for the CNN and 16,643 for the MLP
- Both models reach 1.000 accuracy on the normal validation set
- Shifted accuracy: the score on validation images moved 3 cells down and 3 cells right
- The MLP ties each weight to one fixed pixel position
- Global average pooling lets the CNN answer from which patterns appear, not where they appear
In this example, both models train on the same 240 images with the same batches, optimizer, and schedule. The CNN is the stack from the Training a CNN lesson. The MLP flattens each image to 256 numbers and passes them through a hidden layer of 64. Both models score 1.000 on the normal validation set. The shift then moves every validation image 3 cells down and 3 cells right. The CNN keeps an accuracy of 0.800. The MLP falls to 0.317, no better than guessing among three classes. The table prints the parameter counts as well, 1,299 against 16,643.
Note
On the shifted set the CNN keeps 0.800 and the MLP falls to 0.317, near plain guessing at 0.333.
Try it Yourself
Edit the code, then press Run.
Running is simulated · Python Playground
Exercise
Why does the shifted test hurt the MLP far more than the CNN?