🎓 A1Learner
🔍

PyTorch Training an Ensemble

An ensemble is a group of models whose predictions are combined into one answer. This ensemble averages the class probabilities of three separately trained networks.

Scores become probabilities

Softmax turns raw scores into a row of probabilities. The row sums to 1.0, so rows from different models can be averaged fairly.

scores = torch.tensor([[2.0, 1.0, 0.0]])
probs = torch.softmax(scores, dim=1)
print("Probabilities:", [round(v, 2) for v in probs[0].tolist()])
print("Sum:", round(probs.sum().item(), 2))

Average three rows

Two models lean to class 0 and one leans to class 2. The average follows the two, and the lone model is outvoted.

p1 = torch.tensor([0.7, 0.2, 0.1])
p2 = torch.tensor([0.6, 0.3, 0.1])
p3 = torch.tensor([0.1, 0.3, 0.6])
avg = (p1 + p2 + p3) / 3
print("Averaged:", [round(v, 2) for v in avg.tolist()])
print("Ensemble pick:", avg.argmax().item())

One seed per model

Different seeds give different starting weights. Identical code then trains three genuinely different models.

torch.manual_seed(10)
first = nn.Linear(4, 3)
torch.manual_seed(11)
second = nn.Linear(4, 3)
print("Same start:", torch.equal(first.weight, second.weight))

Example

import torch
from torch import nn

def make_shapes(per_class, seed):
    g = torch.Generator().manual_seed(seed)
    rows = torch.arange(16).reshape(16, 1)
    cols = torch.arange(16).reshape(1, 16)
    images = []
    labels = []
    for label in range(3):
        for _ in range(per_class):
            img = torch.zeros(16, 16)
            if label == 0:
                top = int(torch.randint(6, 9, (1,), generator=g))
                img[top:top + 3, 2:14] = 1.0
            elif label == 1:
                dist = ((rows - 7.5) ** 2 + (cols - 7.5) ** 2).sqrt()
                img[(dist >= 4.0) & (dist <= 5.5)] = 1.0
            else:
                img[4:12, 4:6] = 1.0
                img[4:12, 10:12] = 1.0
                img[4:6, 4:12] = 1.0
                img[10:12, 4:12] = 1.0
            img = img + 0.05 * torch.rand(16, 16, generator=g)
            images.append(img)
            labels.append(label)
    x = torch.stack(images).unsqueeze(1)
    y = torch.tensor(labels)
    perm = torch.randperm(len(labels), generator=g)
    return x[perm], y[perm]


class ShapeCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 8, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(8, 16, 3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
        )
        self.classifier = nn.Linear(16, 3)

    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        return self.classifier(x)


def train_model(seed, x, y, epochs=15):
    torch.manual_seed(seed)
    model = ShapeCNN()
    loss_fn = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    for epoch in range(epochs):
        loss = loss_fn(model(x), y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    return model

x_train, y_train = make_shapes(30, 100)
x_test, y_test = make_shapes(15, 200)

models = [train_model(seed, x_train, y_train) for seed in (10, 11, 12)]

probs = []
for i, model in enumerate(models):
    model.eval()
    with torch.no_grad():
        p = torch.softmax(model(x_test), dim=1)
    probs.append(p)
    acc = (p.argmax(dim=1) == y_test).float().mean().item()
    print(f"Model {i + 1} accuracy:", round(acc, 4))

ensemble_probs = (probs[0] + probs[1] + probs[2]) / 3
ensemble_acc = (ensemble_probs.argmax(dim=1) == y_test).float().mean().item()
print("Ensemble accuracy:", round(ensemble_acc, 4))
print("Test images:", len(y_test))

Key points

  • Ensemble: several models combined into one prediction
  • Use torch.softmax to turn each model's scores into probabilities before averaging
  • Each model trains from its own seed, so the models start differently and err differently
  • Averaging probabilities lets confident models count more than hesitant ones
  • The ensemble is one prediction built from three models, not a fourth trained model
  • Ensemble accuracy uses the same argmax rule as a single model's accuracy

In this example, three copies of the shapes network train on the same 90 images. The network is the canonical shapes CNN, with two convolution blocks and a linear classifier. Seeds 10, 11, and 12 give the three models different starting weights. Each model trains for 15 epochs with Adam at rate 0.01. The test set holds 45 fresh shapes images. The first two models score 1.0, and the third scores 0.6667. The averaged probabilities score 1.0, matching the best single model. Averaging also repaired every test image the third model missed, because its partners were right there.

Note

Softmax rows sum to 1.0 per model, so the averaged row is also a probability row

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

How does this ensemble combine its three models?