🎓 A1Learner
🔍

PyTorch Writing a Custom Layer

A custom layer is a module that wraps new maths in the standard module shape. It stores its learnable numbers in __init__, so an optimizer can find and adjust them.

A layer is a class

The class extends nn.Module and builds its parts in __init__. Temperature is one learnable number that this layer divides its input by. nn.Parameter wraps a plain tensor as that learnable number.

class TemperatureScale(nn.Module):
    def __init__(self):
        super().__init__()
        self.temperature = nn.Parameter(torch.tensor(1.0))

    def forward(self, x):
        return x / self.temperature

Registration puts the number on the list

Assigning an nn.Parameter to self in __init__ registers it. Registered parameters appear in named_parameters() with their names.

layer = TemperatureScale()
for name, param in layer.named_parameters():
    print(name, param.item(), param.requires_grad)
print("Parameter count:", sum(p.numel() for p in layer.parameters()))

The layer divides by its temperature

The temperature is set to 2.0 by hand here. Every score comes out halved.

with torch.no_grad():
    layer.temperature.fill_(2.0)
scores = torch.tensor([4.0, 1.0, -2.0])
print(layer(scores).tolist())

Example

import torch
from torch import nn


class TemperatureScale(nn.Module):
    def __init__(self):
        super().__init__()
        self.temperature = nn.Parameter(torch.tensor(1.0))

    def forward(self, x):
        return x / self.temperature


torch.manual_seed(0)
x = torch.tensor([[2.0, 0.0], [0.0, 2.0], [3.0, 1.0], [1.0, 3.0],
                  [2.5, 0.5], [0.5, 2.5], [3.0, 3.0], [0.0, 0.0]])
y = torch.tensor([0, 1, 0, 1, 0, 1, 0, 1])

model = nn.Sequential(nn.Linear(2, 2), TemperatureScale())
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.05)

with torch.no_grad():
    start_loss = loss_fn(model(x), y).item()
    start_temp = model[1].temperature.item()

for step in range(40):
    loss = loss_fn(model(x), y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

with torch.no_grad():
    final_loss = loss_fn(model(x), y).item()
    final_temp = model[1].temperature.item()
    correct = (model(x).argmax(dim=1) == y).sum().item()

print("Start temperature:", round(start_temp, 3))
print("Start loss:", round(start_loss, 4))
print("Final temperature:", round(final_temp, 3))
print("Final loss:", round(final_loss, 4))
print("Correct:", correct, "of", len(y))

Key points

  • Custom layer: a module class that defines new maths in forward
  • Use nn.Parameter to wrap a tensor as a learnable number on a module
  • Registration: storing a parameter on self inside __init__
  • An optimizer finds registered parameters through model.parameters()
  • A custom layer is a module, so it drops into nn.Sequential like a built-in layer

In this example, the TemperatureScale layer sits after a Linear layer inside nn.Sequential. The data holds eight rows with two numbers each and a class label per row. The temperature starts at 1.0, so the scores pass through unchanged at first. Adam trains the whole model for 40 steps at a learning rate of 0.05. The loss falls from 0.9841 to 0.0029. The temperature itself moves to 0.638, which sharpens every score the Linear layer makes. Dividing by a smaller number stretches the scores apart, so the right class wins by more. The final line shows all eight rows classified correctly.

Note

A plain float stored on self is not a parameter, and no optimizer will adjust it

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Where must a learnable number be stored so an optimizer finds it?