🎓 A1Learner
🔍
0 of 87 lessons completed

PyTorch Training a Sine Predictor

A sine predictor is a model that reads a window of a wave and predicts the next value. This project trains one with an LSTM and checks it on a later stretch of the wave.

Validation sits later in the wave

The first validation target sits at wave index 292. Its value prints as -0.799, far from any training window.

wave = torch.sin(torch.arange(400) * 0.1)
targets = wave[12:]
print(round(targets[280].item(), 3))

Many windows in, one value out

Five windows go in and five values come out. Each window earns exactly one prediction.

torch.manual_seed(0)
model = SinePredictor(16)
batch = torch.zeros(5, 12, 1)
print(tuple(model(batch).shape))

The head reads the last step

The forward pass throws away every step but the last. The linear head turns that one row into the predicted value.

output, _ = self.lstm(x)
return self.head(output[:, -1])

Example

import torch
from torch import nn


def make_windows(wave, width):
    inputs = []
    targets = []
    for start in range(len(wave) - width):
        inputs.append(wave[start:start + width])
        targets.append(wave[start + width])
    return torch.stack(inputs).unsqueeze(-1), torch.stack(targets).unsqueeze(-1)


class SinePredictor(nn.Module):
    def __init__(self, hidden):
        super().__init__()
        self.lstm = nn.LSTM(1, hidden, batch_first=True)
        self.head = nn.Linear(hidden, 1)

    def forward(self, x):
        output, _ = self.lstm(x)
        return self.head(output[:, -1])


wave = torch.sin(torch.arange(400) * 0.1)
x, y = make_windows(wave, 12)
x_train, y_train = x[:280], y[:280]
x_val, y_val = x[280:], y[280:]

torch.manual_seed(0)
model = SinePredictor(16)
loss_fn = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
for epoch in range(50):
    loss = loss_fn(model(x_train), y_train)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

with torch.no_grad():
    val_pred = model(x_val)
    val_error = loss_fn(val_pred, y_val).item()

print("Train windows:", len(x_train))
print("Validation windows:", len(x_val))
print("Validation MSE:", round(val_error, 4))
print("First 10 validation steps (predicted vs true):")
for i in range(10):
    print(round(val_pred[i].item(), 3), "vs", round(y_val[i].item(), 3))

Key points

  • Many-to-one: many input steps produce one output value
  • The head reads the output at the last step only
  • Validation stretch: a later part of the wave that training never used
  • Use nn.MSELoss to score value predictions against the true values
  • One-step-ahead prediction: each validation window predicts its own next value
  • Seed 0 fixes the model's starting weights before training
  • The predictor reads the last step's output, restated plainly: the whole window speaks through one vector

In this example, the wave holds 400 sine values and the window width is 12. The first 280 windows train the model. The remaining 108 windows form the validation stretch, later in the wave. The SinePredictor wraps an LSTM with hidden size 16 and a linear head. Training runs 50 passes with Adam at rate 0.01. The validation MSE prints as 0.002. The pairs then list the first ten validation predictions beside the true values. The predictions trail the true values slightly on the falling stretch, as the first pairs show. By the tenth pair the prediction -0.966 nearly meets the true -0.968. One-step-ahead prediction on a smooth wave is a kind task. The errors lesson earlier in this course shows the habit that matters here. Read the pairs, not just the score.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Why does this project validate on a later stretch of the wave?