🎓 A1Learner
🔍

PyTorch Linear Layers

A linear layer is a module that multiplies its input by a weight and adds a bias. nn.Linear builds one from two numbers, the input size and the output size.

Create a layer

The seed fixes the random starting weight and bias. The weight shape puts the output size first.

torch.manual_seed(0)
layer = nn.Linear(3, 2)
print(layer.weight.shape)
print(layer.bias.shape)

Run one row

One row of three numbers goes in. One row of two numbers comes out.

one = torch.tensor([[1.0, 0.0, 2.0]])
print(layer(one).shape)

Run a batch

A batch is a stack of rows run together. Every row gets the same weight and bias.

batch = torch.tensor([[1.0, 0.0, 2.0],
                      [0.0, 1.0, 1.0],
                      [2.0, 1.0, 0.0],
                      [1.0, 1.0, 1.0]])
print(batch.shape)
print(layer(batch).shape)

Example

import torch
from torch import nn

torch.manual_seed(0)
layer = nn.Linear(3, 2)
batch = torch.tensor([[1.0, 0.0, 2.0],
                      [0.0, 1.0, 1.0],
                      [2.0, 1.0, 0.0],
                      [1.0, 1.0, 1.0]])
output = layer(batch)

print("Input shape:", batch.shape)
print("Weight shape:", layer.weight.shape)
print("Bias shape:", layer.bias.shape)
print("Output shape:", output.shape)
print("Output rows:", output.shape[0])

Key points

  • Use nn.Linear(3, 2) for rows of 3 in and rows of 2 out
  • Weight: the numbers the layer multiplies by, stored as layer.weight
  • Bias: the numbers the layer adds after multiplying, stored as layer.bias
  • Weight shape: (out, in), so nn.Linear(3, 2) has weight shape (2, 3)
  • A linear layer is one multiply and one add, shared by every row

In this example, the seed is set first, so the starting weight and bias are repeatable. The layer is built as nn.Linear(3, 2). The batch holds four rows of three numbers. Calling the layer runs the batch through it. The prints show the input shape first. The weight shape puts the output size first, as (2, 3). The bias holds one number per output. The output shape shows four rows of two numbers.

Note

Linear layers are also called dense or fully connected layers in other tools

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

nn.Linear(3, 2) expects each input row to hold how many numbers?