Skip to course content
Free PyTorch course

Advanced Deep Learning with PyTorch

Unit 09.02: Using pretrained models for feature extraction

The cheapest useful thing you can do with a pretrained model: run it once, keep the vectors, and fit something small on top.

Embed once, reuse many times

The encoder is frozen and switched to eval mode. Every document is passed through once, pooled to a single vector, and those vectors are cached.

From there a logistic regression — or any classical model — trains in milliseconds:

import torch
from torch import nn
from sklearn.linear_model import LogisticRegression

torch.manual_seed(0)
# Stand-in for a pretrained encoder: frozen, used only to produce features.
encoder = nn.TransformerEncoderLayer(d_model=16, nhead=4, dim_feedforward=32,
                                     batch_first=True)
for p in encoder.parameters():
    p.requires_grad = False
encoder.eval()

tokens = torch.randn(120, 6, 16)
labels = (tokens[:, 0, 0] > 0).long()

with torch.no_grad():                       # no gradients needed at all
    features = encoder(tokens).mean(dim=1)  # mean-pool over tokens

print("features:", tuple(features.shape), "one vector per document")

clf = LogisticRegression(max_iter=1000).fit(features[:90].numpy(), labels[:90].numpy())
acc = clf.score(features[90:].numpy(), labels[90:].numpy())
print("logistic head on frozen features:", round(float(acc), 3))
print("trainable transformer parameters :",
      sum(p.numel() for p in encoder.parameters() if p.requires_grad))

# This is the cheapest useful thing you can do with a pretrained model: embed
# once, cache the vectors, and fit a small classical model on top.

Zero trainable transformer parameters, and a working classifier. The entire cost was one forward pass per document, done inside torch.no_grad() so no graph was built.

This approach scales well operationally: the expensive step happens once, and downstream experiments are cheap enough to iterate on freely.

The mistake this prevents

Re-encoding the same documents for every experiment. The embedding pass dominates the runtime, and caching it turns a slow loop into a fast one.

Takeaway

Freeze, embed once, cache the vectors, and iterate on a cheap head.