PyTorch Advanced Capstone

torch.jit.load rebuilds the model from the file alone. The ReviewClassifier class is never imported at load time.

print(loaded(test_x[:2]).shape)

The classifier is small by design. Its parameter total fits in a phone-sized file with room to spare.

total = sum(p.numel() for p in model.parameters())
print(f"Parameters: {total}")

The byte count from the closing example reads as kilobytes with one division. The scripted file holds code and weights together.

print(f"Scripted file size: {os.path.getsize('capstone_model.pt') / 1024:.1f} KB")

import os
import random

import torch
from torch import nn

POSITIVE = ["great", "lovely", "bright", "fun", "warm"]
NEGATIVE = ["awful", "dull", "cold", "boring", "grim"]
FILLER = ["the", "a", "movie", "story", "scene", "really", "very"]
VOCAB = POSITIVE + NEGATIVE + FILLER + ["not"]
WORD_ID = {word: i for i, word in enumerate(VOCAB)}
LENGTH = 5


def make_data(rng, n_plain, n_negated):
    data = []
    for _ in range(n_plain):
        label = rng.randrange(2)
        side = POSITIVE if label == 1 else NEGATIVE
        words = [rng.choice(side) for _ in range(3)]
        words += [rng.choice(FILLER) for _ in range(LENGTH - 3)]
        rng.shuffle(words)
        data.append((words, label))
    for _ in range(n_negated):
        side_label = rng.randrange(2)
        side = POSITIVE if side_label == 1 else NEGATIVE
        rest = [rng.choice(side) for _ in range(2)]
        rest += [rng.choice(FILLER) for _ in range(LENGTH - 3)]
        rng.shuffle(rest)
        data.append((["not"] + rest, 1 - side_label))
    rng.shuffle(data)
    return data


def encode(data):
    x = torch.tensor([[WORD_ID[w] for w in words] for words, _ in data])
    y = torch.tensor([label for _, label in data])
    return x, y


class ReviewClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.embedding = nn.Embedding(len(VOCAB), 16)
        self.position = nn.Embedding(LENGTH, 16)
        layer = nn.TransformerEncoderLayer(
            d_model=16, nhead=2, dim_feedforward=32,
            dropout=0.0, batch_first=True)
        self.encoder = nn.TransformerEncoder(layer, num_layers=1)
        self.head = nn.Linear(16, 2)

    def forward(self, ids):
        h = self.encoder(self.embedding(ids) + self.position.weight)
        return self.head(h.mean(dim=1))


rng = random.Random(42)
train_x, train_y = encode(make_data(rng, 96, 64))

torch.manual_seed(0)
model = ReviewClassifier()
optimizer = torch.optim.Adam(model.parameters(), lr=0.02)
loss_fn = nn.CrossEntropyLoss()

for step in range(60):
    optimizer.zero_grad()
    loss = loss_fn(model(train_x), train_y)
    loss.backward()
    optimizer.step()

print(f"Final training loss: {float(loss):.4f}")

test_rng = random.Random(7)
test_x, test_y = encode(make_data(test_rng, 120, 80))

model.eval()
with torch.no_grad():
    predictions = model(test_x).argmax(dim=1)
accuracy = float((predictions == test_y).float().mean())
print(f"Held-out accuracy: {accuracy:.3f}")

scripted = torch.jit.script(model)
scripted.save("capstone_model.pt")
size = os.path.getsize("capstone_model.pt")
print(f"Scripted file size: {size} bytes")

loaded = torch.jit.load("capstone_model.pt")
loaded.eval()
with torch.no_grad():
    loaded_predictions = loaded(test_x).argmax(dim=1)
loaded_accuracy = float((loaded_predictions == test_y).float().mean())
print(f"Loaded model accuracy: {loaded_accuracy:.3f}")
print(f"Predictions match: {bool((loaded_predictions == predictions).all())}")