🎓 A1Learner
🔍

PyTorch nn.Embedding

nn.Embedding is a lookup table of word vectors that training can shape. An index goes in, and the row stored for that index comes out.

Build the table

The table for 10 words at width 4 is a (10, 4) tensor. requires_grad is True, so training can shape the rows.

embedding = nn.Embedding(10, 4)
print(tuple(embedding.weight.shape))
print(embedding.weight.requires_grad)

Indices in, rows out

Two indices go in as one tensor. Two rows come out, stacked in the same order.

picked = embedding(torch.tensor([2, 5]))
print(picked.shape)

The lookup copies the row

A lookup does not compute a new vector. It returns the stored row for that index, exactly as the table holds it.

row = embedding(torch.tensor(3)).detach()
print(torch.equal(row, embedding.weight[3].detach()))

Example

import torch
from torch import nn

torch.manual_seed(3)

words = ["great", "lovely", "bright", "fun", "warm",
         "awful", "dull", "cold", "boring", "grim"]
embedding = nn.Embedding(len(words), 4)

print("Table shape:", tuple(embedding.weight.shape))
picked = embedding(torch.tensor([0, 4, 6])).detach()
for word, row in zip(["great", "warm", "dull"], picked):
    print(word, "->", row)
again = embedding(torch.tensor([0])).detach()
print("Index 0 gives the same row twice:", torch.equal(picked[0], again[0]))

Key points

  • Embedding layer: a table whose rows are word vectors
  • Use nn.Embedding(vocab_size, dim) to build the table
  • Input and output: a tensor of indices goes in, and the matching rows come out stacked
  • The same index always returns the same row
  • embedding.weight: the table itself, held as a trainable parameter

In this example, the seed is set first, so the starting table is the same on every run. The table holds 10 words at width 4. The lookup asks for indices 0, 4, and 6, the words great, warm, and dull. Each word prints with its row of four numbers. A second lookup of index 0 returns the same row, and the equality check prints True. The rows are still random at this point. Training is what moves them into useful places, as the Bag of Embeddings lesson shows.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does nn.Embedding return for a tensor of three indices?