N-BEATS › Implementation in PyTorch
Machine Learning from ScratchAll chapters N-BEATS › Implementation in PyTorchOne section · Sections in this chapter5 / 8 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Now that we understand the mathematical framework and have seen a working example, let's build a complete, production-ready implementation. This implementation will translate the formulas we've discussed into actual PyTorch code that you can use for your own time series forecasting projects.
What we'll implement: We'll create a full N-BEATS model that includes:
- The neural network architecture that learns coefficients ()
- The trend basis functions ( with polynomial terms)
- The seasonal basis functions ( with trigonometric terms)
- The training loop that minimizes the loss function ()
This implementation follows the mathematical framework we've established, making it easy to see how the formulas translate into code.
Step 1: Import Required Libraries
Step 2: Define the N-BEATS Architecture
class NBEATSBlock(nn.Module):
"""Single N-BEATS block for trend or seasonal decomposition"""
def __init__(self, input_size, theta_size, hidden_size=64):
super(NBEATSBlock, self).__init__()
# Neural network to learn coefficients
self.network = nn.Sequential(
nn.Linear(input_size, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, hidden_size),
nn.ReLU(),
nn.Linear(hidden_size, theta_size),
)
self.theta_size = theta_size
def forward(self, x):
return self.network(x)
class NBEATS(nn.Module):
"""Complete N-BEATS model"""
def __init__(self, input_size=12, forecast_horizon=6, hidden_size=64):
super(NBEATS, self).__init__()
# Trend block
self.trend_block = NBEATSBlock(
input_size=input_size,
theta_size=forecast_horizon
* 2, # Linear and quadratic coefficients
hidden_size=hidden_size,
)
# Seasonal block
self.seasonal_block = NBEATSBlock(
input_size=input_size,
theta_size=forecast_horizon * 2, # Sine and cosine coefficients
hidden_size=hidden_size,
)
self.forecast_horizon = forecast_horizon
def forward(self, x):
# Get trend coefficients
trend_coeffs = self.trend_block(x)
trend_coeffs = trend_coeffs.view(-1, self.forecast_horizon, 2)
# Get seasonal coefficients
seasonal_coeffs = self.seasonal_block(x)
seasonal_coeffs = seasonal_coeffs.view(-1, self.forecast_horizon, 2)
# Generate forecasts
forecasts = []
for h in range(self.forecast_horizon):
# Trend component
trend_h = trend_coeffs[:, h, 0] * (h + 1) + trend_coeffs[
:, h, 1
] * ((h + 1) ** 2)
# Seasonal component
seasonal_h = seasonal_coeffs[:, h, 0] * torch.sin(
2 * torch.pi * torch.tensor(h + 1) / 12
) + seasonal_coeffs[:, h, 1] * torch.cos(
2 * torch.pi * torch.tensor(h + 1) / 12
)
forecast_h = trend_h + seasonal_h
forecasts.append(forecast_h)
return torch.stack(forecasts, dim=1)
# Initialize model
model = NBEATS(input_size=12, forecast_horizon=6, hidden_size=64)
total_params = sum(p.numel() for p in model.parameters())Model parameters: 11544
The model contains a moderate number of parameters, which balances model capacity with computational efficiency. For this configuration with input_size=12, forecast_horizon=6, and hidden_size=64, the parameter count is reasonable for most time series forecasting tasks. This size allows the model to learn complex patterns without requiring excessive computational resources or risking overfitting on smaller datasets.
Step 3: Data Preparation and Training
def prepare_data(data, input_size=12, output_size=6):
"""Prepare data for N-BEATS training"""
X, y = [], []
for i in range(len(data) - input_size - output_size + 1):
X.append(data[i : i + input_size])
y.append(data[i + input_size : i + input_size + output_size])
return np.array(X), np.array(y)
# Generate sample data
np.random.seed(42)
months = 100
t = np.arange(months)
# Create time series with trend and seasonality
trend = 100 + 0.5 * t
seasonal = 10 * np.sin(2 * np.pi * t / 12) + 5 * np.cos(2 * np.pi * t / 12)
noise = np.random.normal(0, 2, months)
data = trend + seasonal + noise
# Prepare training data
X, y = prepare_data(data, input_size=12, output_size=6)
# Convert to PyTorch tensors
X_tensor = torch.FloatTensor(X)
y_tensor = torch.FloatTensor(y)Training data shape: torch.Size([83, 12]) Target data shape: torch.Size([83, 6]) Number of training samples: 83
The training data consists of input sequences of length 12 (one year of monthly data) and target sequences of length 6 (6-month forecasts). The number of training samples is determined by the sliding window approach, which creates overlapping sequences from the historical data. With 100 months of data, the sliding window creates multiple training examples, allowing the model to learn from different segments of the time series. This approach maximizes the use of available data and helps the model generalize better.
Step 4: Training Loop
# Training setup
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.MSELoss()
# Training loop
epochs = 1000
model.train()
training_losses = []
for epoch in range(epochs):
optimizer.zero_grad()
# Forward pass
predictions = model(X_tensor)
# Calculate loss
loss = criterion(predictions, y_tensor)
# Backward pass
loss.backward()
optimizer.step()
training_losses.append(loss.item())Epoch 0, Loss: 19513.5293 Epoch 200, Loss: 40.3587 Epoch 400, Loss: 19.7800 Epoch 600, Loss: 10.7850 Epoch 800, Loss: 7.3420 Final training loss: 6.0384
The training loss decreases steadily over epochs, indicating that the model is successfully learning to capture the trend and seasonal patterns in the data. The steady decline suggests the model is converging rather than overfitting, as the loss continues to decrease without signs of instability. The final loss value provides a baseline for comparison. Lower values indicate better fit to the training data, though you should also evaluate on validation data to assess generalization performance.
Step 5: Making Predictions
# Make prediction on last window
model.eval()
with torch.no_grad():
last_window = data[-12:].reshape(1, -1)
last_window_tensor = torch.FloatTensor(last_window)
forecast = model(last_window_tensor)
forecast_np = forecast.numpy().flatten()
# Generate "actual" future values for comparison (using the same generating process)
np.random.seed(43) # Different seed for future data
t_future = np.arange(months, months + 6)
trend_future = 100 + 0.5 * t_future
seasonal_future = 10 * np.sin(2 * np.pi * t_future / 12) + 5 * np.cos(
2 * np.pi * t_future / 12
)
noise_future = np.random.normal(0, 2, 6)
actual_future = trend_future + seasonal_future + noise_future
# Calculate performance metrics
mae = np.mean(np.abs(forecast_np - actual_future))
rmse = np.sqrt(np.mean((forecast_np - actual_future) ** 2))N-BEATS Forecast: Month 1: 156.45 Month 2: 151.93 Month 3: 147.65 Month 4: 141.39 Month 5: 140.69 Month 6: 144.73 Performance Metrics: MAE: 1.74 RMSE: 2.06
The forecast shows the predicted values for the next 6 months, extrapolating the learned trend and seasonal patterns from the training data. The MAE (Mean Absolute Error) measures the average absolute difference between predictions and actual values, while RMSE (Root Mean Squared Error) penalizes larger errors more heavily. Lower values for both metrics indicate better forecast accuracy. For this synthetic dataset, the relatively low MAE and RMSE values suggest the model has successfully learned the underlying trend and seasonal patterns. In practice, compare these metrics to baseline methods (such as naive forecasts or seasonal averages) to assess whether N-BEATS provides meaningful improvements.
Key Parameters
Below are some of the main parameters that affect how N-BEATS works and performs.
-
input_size: Number of historical time series values used as input (default: 12). Should match or exceed the seasonal period. Use 12 for monthly data with yearly seasonality, 7 or 14 for daily data with weekly patterns. Larger windows capture more historical context but may include irrelevant distant patterns. -
forecast_horizon: Number of future periods to forecast (default: 6). Shorter horizons (1-6 periods) are typically more accurate but may not provide sufficient lead time. Longer horizons (12-24 periods) are useful for strategic planning but require more historical data. Keep the forecast horizon to at most half the length of your training data. -
hidden_size: Size of hidden layers in the neural network (default: 64). Controls model complexity and capacity. Larger values (128-256) can capture more complex patterns but increase computation time and risk overfitting with limited data. Smaller values (32-64) work well for simpler patterns and smaller datasets. -
learning_rate: Learning rate for the Adam optimizer (default: 0.001). Controls how quickly the model learns. Lower values (0.0001-0.001) provide more stable training but may require more epochs. Higher values (0.01-0.1) train faster but may be unstable. Start with 0.001 and adjust based on training behavior. -
epochs: Number of training iterations (default: 1000). More epochs allow the model to learn better but increase training time. Monitor training and validation loss to determine when to stop. Use early stopping if validation loss stops improving. -
random_state: Seed for reproducibility (default: None). Set to an integer to ensure consistent results across runs. Important for debugging and comparing different model configurations.
Key Methods
The following are the most commonly used methods for interacting with N-BEATS models.
-
forward(x): Generates forecasts for input sequences. Takes a tensor of shape(batch_size, input_size)containing historical time series values and returns predictions of shape(batch_size, forecast_horizon). This is called automatically when usingmodel(x). -
train(): Sets the model to training mode, enabling dropout and batch normalization updates. Call this before training to ensure the model learns from the data. -
eval(): Sets the model to evaluation mode, disabling dropout and freezing batch normalization statistics. Call this before making predictions to ensure consistent, deterministic outputs. -
parameters(): Returns an iterator over all model parameters (weights and biases). Used by optimizers to update parameters during training. Can also be used to inspect or modify specific parameters.