Traffic Flow Prediction with GCN-GRU
Build a spatio-temporal deep learning model to predict LA highway traffic speeds.
Introduction
30 Second Summary
Every driver knows the frustration of hitting unexpected traffic on a highway. Navigation apps predict congestion using speed data from road sensors, but the real challenge is that traffic at one sensor depends on what is happening at nearby sensors.
In this project, you will build a deep learning model that predicts traffic speeds across 207 real highway sensors in Los Angeles. You will compare a temporal-only baseline against a model that understands the road network's spatial structure, proving that graph neural networks outperform traditional approaches on networked sensor data.
What You'll Build
Run your finished notebook and see the model predict the next 5 minutes of traffic speed across all 207 LA highway sensors, with side-by-side plots proving the graph-powered model wins.
By the end of this project, you'll have:
- A Google Colab notebook that loads real LA highway traffic data and visualizes the 207-sensor road network as a color-coded graph.
- A trained GRU-only baseline model with prediction plots that reveal where ignoring spatial relationships causes predictions to lag behind sudden speed changes.
- A trained T-GCN model that fuses Graph Convolutional Networks with Gated Recurrent Units, with comparison charts and a per-sensor heatmap proving the spatial model's accuracy advantage.
- Secret Mission: Extend your model to predict 15, 30, and 60 minutes into the future and chart how prediction accuracy degrades at longer horizons.
Are there any prerequisites?
You need a free Google account to use Google Colab. No prior deep learning or graph neural network experience is needed.
Before We Start
Before diving into code, let's lock in what you're about to build and why it matters.
Set Up Google Colab and Load the METR-LA Dataset
Training a deep learning model requires two things up front: a GPU to handle the heavy matrix math, and real-world data to learn from. Google Colab gives you both for free.
In this step, you will open a Colab notebook with GPU acceleration, install PyTorch Geometric (the graph neural network library), and download the METR-LA traffic speed dataset. By the end, every sensor reading and road connection will be loaded in memory and ready for modeling.
In this step, get ready to:
- Open Google Colab with a T4 GPU runtime.
- Install PyTorch Geometric and download the METR-LA dataset from Zenodo.
- Load the speed matrix and adjacency matrix into memory and verify their shapes.
Open Google Colab and enable a GPU runtime
Google Colab is a free browser-based notebook environment that comes with PyTorch pre-installed and access to a GPU. You do not need to install anything on your own computer.
- Open colab.research.google.com in your browser and sign in with your Google account.
- Click New Notebook in the bottom-right corner of the pop-up window. If you do not see a pop-up, click File in the top menu bar, then New notebook.
- Click Runtime in the top menu bar.
- Click Change runtime type.
- In the pop-up, set Hardware accelerator to T4 GPU.
- Click Save.
What is a GPU runtime?
A GPU (Graphics Processing Unit) is a chip designed for parallel math. Deep learning models multiply massive matrices during training, and a GPU can process these operations hundreds of times faster than a regular CPU.
Colab's free tier provides access to an NVIDIA T4 GPU. Sessions last up to 12 hours and GPU availability varies by demand.
Install PyTorch Geometric and download the METR-LA dataset
Colab already ships with PyTorch, but you need PyTorch Geometric for graph neural network layers. You also need the METR-LA dataset, which is a public benchmark of traffic speeds from 207 highway sensors in Los Angeles.
- Install PyTorch Geometric, import all the libraries you will use across the project, and verify your environment by pasting this code into the first cell and running it:
# Install PyTorch Geometric (Colab already has PyTorch)
!pip install torch_geometric -q
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import pickle
import os
import urllib.request
from torch_geometric.nn import GCNConv
from torch_geometric.utils import dense_to_sparse
# Verify environment
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")
What does this code do?
- !pip install torch_geometric -q installs PyTorch Geometric quietly. The ! prefix runs shell commands directly from a Colab cell.
- The GCNConv import gives you the Graph Convolutional Network layer you will use later to capture spatial relationships between sensors.
- dense_to_sparse is a utility that converts a dense adjacency matrix into the edge list format PyTorch Geometric expects.
- The three print lines confirm your PyTorch version, whether a GPU is available, and which device (CPU or CUDA) the notebook will use.
- Run the cell by pressing Shift+Enter.
- You should see pip install output followed by your PyTorch version and CUDA available: True.
Seeing CUDA available: False?
- Go back to Runtime > Change runtime type and confirm T4 GPU is selected. Click Save, then re-run the cell.
- GPU availability fluctuates on the free tier. If no GPU is available, wait a few minutes and try again, or try a different browser session.
- Still stuck? Help me troubleshoot CUDA not being available in Google Colab.
- Click + Code in the top-left toolbar to create a new code cell.
- Download the METR-LA speed data and adjacency matrix from Zenodo by pasting this code into the new cell:
# Download METR-LA dataset from Zenodo
os.makedirs('data', exist_ok=True)
metr_la_url = "https://zenodo.org/records/5724362/files/METR-LA.csv?download=1"
adj_url = "https://zenodo.org/records/5724362/files/adj_mx_METR-LA.pkl?download=1"
if not os.path.exists('data/METR-LA.csv'):
print("Downloading METR-LA.csv...")
urllib.request.urlretrieve(metr_la_url, 'data/METR-LA.csv')
if not os.path.exists('data/adj_mx_METR-LA.pkl'):
print("Downloading adj_mx_METR-LA.pkl...")
urllib.request.urlretrieve(adj_url, 'data/adj_mx_METR-LA.pkl')
What does this code do?
- os.makedirs('data', exist_ok=True) creates a data/ folder in the Colab runtime to store the downloaded files.
- METR-LA.csv contains traffic speed readings from 207 highway loop detectors in Los Angeles, recorded every 5 minutes from March to June 2012.
- adj_mx_METR-LA.pkl contains the adjacency matrix. This matrix encodes which sensors are connected to each other on the road network.
- The if not os.path.exists checks skip the download if the files already exist, so you can safely re-run the cell.
- Run the cell. You should see Downloading METR-LA.csv... and Downloading adj_mx_METR-LA.pkl... printed as the files download.
Where does this data come from?
The METR-LA dataset is hosted on Zenodo under a CC BY 4.0 license, which means it is free to use. It is one of the most widely used benchmarks for traffic prediction research.
Download taking too long or failing?
- The CSV file is about 72 MB. On a slow connection it may take a minute or two. Wait for the cell to finish before moving on.
- If you see a connection error, re-run the cell. The if not os.path.exists check means it will only re-download the file that failed.
- Help me fix a download error when fetching the METR-LA dataset in Google Colab.
Now that the files are on disk, you need to load them into Python variables so you can work with them.
Before you run the next cell, think about what shape you expect the speed matrix to have. There are 207 sensors recording at 5-minute intervals over roughly 4 months. How many rows does that produce?
- Click + Code to create another new cell.
- Load the speed data and adjacency matrix into memory by pasting this code:
# Load speed data
df = pd.read_csv('data/METR-LA.csv', index_col=0)
speed_matrix = df.values.astype(np.float32)
print(f"Speed matrix shape: {speed_matrix.shape}")
# Load adjacency matrix
with open('data/adj_mx_METR-LA.pkl', 'rb') as f:
sensor_ids, sensor_id_to_ind, adj_mx = pickle.load(f, encoding='latin1')
adj_mx = adj_mx.astype(np.float32)
print(f"Adjacency matrix shape: {adj_mx.shape}")
print(f"Number of sensors: {len(sensor_ids)}")
print(f"Date range: {df.index[0]} to {df.index[-1]}")
What does this code do?
- pd.read_csv loads the CSV into a pandas DataFrame. The index_col=0 parameter treats the first column (timestamps) as the row index.
- speed_matrix is the raw NumPy array of traffic speeds. Each row is one 5-minute timestep. Each column is one sensor.
- pickle.load unpacks the adjacency matrix file into three objects: a list of sensor IDs, a mapping from sensor ID to index, and the 207x207 adjacency matrix itself.
- adj_mx tells your model which sensors are neighbors on the road network. A non-zero value at position (i, j) means sensor i and sensor j are connected.
- Run the cell. You should see the speed matrix shape printed as (34272, 207) and the adjacency matrix shape as (207, 207).
That is 34,272 five-minute timesteps across 207 sensors. The adjacency matrix is a 207-by-207 grid encoding which sensors are physically connected on the highway network.
Seeing a FileNotFoundError?
- This means the download cell did not finish successfully. Go back to the previous cell and re-run it. Wait for it to complete before running this cell again.
- Make sure you ran the cells in order. Colab cells must be run top to bottom.
- Help me fix a FileNotFoundError when loading the METR-LA data.
✔️ Awesome, I've got everything!
Great. Make sure all three cells ran without errors before moving on.
ⓧ I'd like to double check the full code
Scroll up through your notebook and confirm you have all three cells, in this order:
- Cell 1a: starts with !pip install torch_geometric -q, contains all imports, and ends with the three print lines verifying PyTorch and CUDA.
- Cell 1b: starts with os.makedirs('data', exist_ok=True) and downloads both files from Zenodo.
- Cell 1c: loads METR-LA.csv into speed_matrix and adj_mx_METR-LA.pkl into adj_mx, then prints both shapes.
Each cell should show output underneath it. The final cell's output should include Speed matrix shape: (34272, 207) and Adjacency matrix shape: (207, 207).
Your environment is ready and the data is loaded. Next up, you will visualize the sensor network as a graph and explore how traffic patterns flow across it.
Explore the Traffic Graph and Visualize Patterns
Your data is loaded and your environment is ready. Now you need to understand what this data actually looks like before building any models.
The METR-LA dataset captures traffic speeds across a real highway sensor network. Each sensor is connected to nearby sensors on the road, forming a graph. In this step, you will visualize that graph, spot patterns in the traffic speeds, and prepare structured training data that your deep learning models can learn from.
In this step, get ready to:
- Convert the adjacency matrix into graph edges and visualize the 207-sensor road network.
- Plot traffic speed time series to see daily rush-hour patterns and spatial correlations.
- Prepare normalized training, validation, and test splits using sliding windows.
Convert the adjacency matrix to graph edges
The adjacency matrix you loaded is a 207x207 grid where each value represents the connection strength between two sensors. PyTorch Geometric needs this in a different format: a list of edges (pairs of connected nodes) and their weights.
- Create a new code cell in your notebook and add the following code to convert the adjacency matrix and print graph statistics:
import networkx as nx
# Convert adjacency to edge_index for PyTorch Geometric
adj_tensor = torch.tensor(adj_mx, dtype=torch.float32)
edge_index, edge_weight = dense_to_sparse(adj_tensor)
# Remove self-loops for cleaner visualization
mask = edge_index[0] != edge_index[1]
edge_index_no_self = edge_index[:, mask]
edge_weight_no_self = edge_weight[mask]
# Filter to meaningful edges (weight > 0.1)
sig_mask = edge_weight_no_self > 0.1
edge_index_vis = edge_index_no_self[:, sig_mask]
print(f"Number of nodes: {adj_mx.shape[0]}")
print(f"Number of edges (weight > 0.1, no self-loops): {edge_index_vis.shape[1]}")
print(f"Avg connections per sensor: {edge_index_vis.shape[1] / adj_mx.shape[0]:.1f}")
What does this code do?
dense_to_sparse converts the dense 207x207 matrix into two tensors: edge_index (pairs of connected nodes) and edge_weight (how strongly connected each pair is).
Self-loops (a sensor connected to itself) are removed because they add no spatial information. Edges with weight below 0.1 are filtered out to keep only meaningful connections between nearby sensors.
- Run the cell.
You should see the number of nodes printed as 207, followed by the edge count and average connections per sensor.
Seeing a NameError?
Make sure you ran Cell 1 (from Step 1) first. The variables adj_mx and dense_to_sparse are defined there. If your Colab runtime disconnected, re-run Cell 1 before this cell.
Seeing a NameError for adj_mx or dense_to_sparse?
Visualize the sensor network and traffic patterns
Numbers alone do not tell you much about a road network. A visual of the graph reveals clusters of sensors and how densely interconnected different parts of the highway system are. Plotting individual sensor speeds over time reveals daily rush-hour dips and shows that nearby sensors behave similarly.
- Add this code in a new cell to build a NetworkX graph and plot both the sensor network and traffic speed time series:
# Build networkx graph for visualization
G = nx.Graph()
for i in range(adj_mx.shape[0]):
G.add_node(i)
for k in range(edge_index_vis.shape[1]):
src, dst = edge_index_vis[0, k].item(), edge_index_vis[1, k].item()
G.add_edge(src, dst)
avg_speed = np.mean(speed_matrix, axis=0)
fig, axes = plt.subplots(1, 2, figsize=(16, 6))
# Plot 1: Sensor network graph
pos = nx.spring_layout(G, seed=42, k=0.3)
nx.draw_networkx(G, pos, ax=axes[0], node_size=20, with_labels=False,
node_color=avg_speed, cmap='RdYlGn', edge_color='lightgray',
width=0.3, alpha=0.8)
axes[0].set_title('METR-LA Sensor Network (207 sensors)\nColor = Average Speed')
# Plot 2: Speed time series for 3 sensors
sensors_to_plot = [0, 50, 150]
timesteps = 2016 # 7 days
hours = np.arange(timesteps) * 5 / 60 # Convert to hours
for s in sensors_to_plot:
axes[1].plot(hours, speed_matrix[:timesteps, s], label=f'Sensor {s}', alpha=0.7)
axes[1].set_xlabel('Hours')
axes[1].set_ylabel('Speed (mph)')
axes[1].set_title('Traffic Speed - First 7 Days')
axes[1].legend()
axes[1].grid(True, alpha=0.3)
plt.tight_layout()
plt.show()
What's happening here?
The left plot draws the 207 sensors as nodes with edges showing road connections. Node color maps to average speed, so green nodes are faster stretches and red nodes are slower or more congested areas.
The right plot shows raw speed readings for sensors 0, 50, and 150 over 7 days. Each dip is a rush-hour slowdown. Notice how the three sensors dip at roughly the same times. That correlation between nearby sensors is the spatial signal that a Graph Convolutional Network can capture.
- Run the cell.
You should see two side-by-side plots. On the left, a network graph of 207 nodes with gray edges connecting them, colored from red to green by average speed. On the right, three overlapping speed lines that dip together during rush hours and recover overnight.
Graph looks like a single blob?
The spring layout spreads nodes based on connections, so a densely connected graph can appear clumped. This is expected for a highway network where many sensors are closely linked. The important thing is that you see distinct nodes and edges, not a blank canvas.
Why does my sensor network graph look so dense?
Prepare the training data
Deep learning models cannot learn from raw speed readings directly. You need to structure the data into input-output pairs: given the last hour of speeds (12 five-minute timesteps), predict the next 5 minutes. You also need to normalize the speeds so the model trains efficiently.
- Add this code in a new cell to define the normalization and sliding window function:
# Prepare training data
SEQ_LEN = 12 # 1 hour of input (12 x 5min)
PRED_LEN = 1 # Predict next 5 minutes
NUM_NODES = speed_matrix.shape[1]
# Z-score normalization
mean_speed = np.mean(speed_matrix)
std_speed = np.std(speed_matrix)
speed_normalized = (speed_matrix - mean_speed) / std_speed
# Create sequences
def create_sequences(data, seq_len, pred_len):
X, Y = [], []
for i in range(len(data) - seq_len - pred_len + 1):
X.append(data[i:i+seq_len])
Y.append(data[i+seq_len:i+seq_len+pred_len])
return np.array(X), np.array(Y)
What does this code do?
- SEQ_LEN = 12 means each input sample covers 12 timesteps (1 hour of 5-minute readings).
- Z-score normalization subtracts the mean and divides by the standard deviation, centering all speeds around 0. This helps gradient-based optimizers converge faster.
- create_sequences slides a window across the timeline, collecting 12-timestep inputs and 1-timestep targets. Each window shifts forward by one step.
- Run the cell.
- Now create the actual training splits by adding this code in a new cell:
X, Y = create_sequences(speed_normalized, SEQ_LEN, PRED_LEN)
Y = Y.squeeze(1) # Shape: (samples, num_nodes)
# Chronological split: 70% train, 10% val, 20% test
n = len(X)
train_end = int(n * 0.7)
val_end = int(n * 0.8)
X_train, Y_train = X[:train_end], Y[:train_end]
X_val, Y_val = X[train_end:val_end], Y[train_end:val_end]
X_test, Y_test = X[val_end:], Y[val_end:]
print(f"\nDataset splits:")
print(f"Train: {X_train.shape[0]} samples")
print(f"Val: {X_val.shape[0]} samples")
print(f"Test: {X_test.shape[0]} samples")
print(f"Input shape: (batch, {SEQ_LEN} timesteps, {NUM_NODES} sensors)")
print(f"Target shape: (batch, {NUM_NODES} sensors)")
Why a chronological split?
Traffic data is ordered in time, so you cannot shuffle it randomly. If you did, the model could "peek" at future data during training. Instead, the first 70% of timesteps are used for training, the next 10% for validation (tuning), and the final 20% for testing.
Before you run this cell, how many training samples do you think you will get from 34,272 timesteps with a window of 12?
- Run the cell.
You should see the dataset split summary printed. The train set has roughly 24,000 samples, the validation set around 3,400, and the test set around 6,800. Each input has the shape (batch, 12 timesteps, 207 sensors) and each target has the shape (batch, 207 sensors).
Seeing a MemoryError or kernel crash?
The create_sequences function builds large NumPy arrays. If your Colab runtime is low on memory, go to Runtime then Restart runtime and re-run all cells from the top.
Why is my Colab kernel crashing when creating sequences?
✔️ Awesome, I've got everything!
Make sure you have saved your notebook. You now have four cells of code across Steps 1 and 2.
ⓧ I'd like to double check the full code
Scroll through the cells you added in this step and confirm each chunk is present, in order:
- Cell with import networkx as nx at the top. Converts the adjacency matrix, filters edges, and prints node/edge counts.
- Cell that builds G = nx.Graph(), computes avg_speed, and creates the two side-by-side plots (sensor network and 7-day speed time series).
- Cell that defines SEQ_LEN, PRED_LEN, NUM_NODES, normalizes the speed matrix, and defines the create_sequences function.
- Cell that calls create_sequences, splits into train/val/test, and prints the dataset split summary.
You have explored the graph, seen the spatial correlations in action, and prepared your data. Next up, you will build a GRU-only baseline model that ignores all of this spatial structure and treats each sensor independently.
Build and Train a GRU-Only Baseline
In Step 2, you explored the sensor network and saw that nearby sensors show correlated traffic patterns. That spatial signal is exactly what a graph-aware model can exploit.
Before building that graph-aware model, you need a baseline to beat. In this step, you will build a GRU-only model that processes each sensor's time series independently, completely ignoring the road network. Seeing where this model falls short will make the case for adding spatial awareness in the next step.
In this step, get ready to:
- Define a GRU baseline model that treats each sensor independently.
- Train the model for 50 epochs and evaluate its test MAE.
- Visualize predictions vs actuals to identify the baseline's limitations.
Define the GRU baseline model
A Gated Recurrent Unit (GRU) is a type of recurrent neural network that learns patterns from sequential data. It processes one timestep at a time, carrying a hidden state forward through the sequence. For traffic prediction, this means it can learn daily rush-hour dips and recovery patterns.
The key limitation: this baseline reshapes the input so each sensor is processed as its own independent time series. Sensor 0 has no idea what Sensor 1 is reporting, even if they sit on the same highway.
- Create a new code cell below Cell 2 and add the GRUBaseline class by pasting this code:
class GRUBaseline(nn.Module):
def __init__(self, num_nodes, input_dim=1, hidden_dim=32):
super(GRUBaseline, self).__init__()
self.num_nodes = num_nodes
self.hidden_dim = hidden_dim
self.gru = nn.GRU(input_size=input_dim, hidden_size=hidden_dim, batch_first=True)
self.fc = nn.Linear(hidden_dim, 1)
def forward(self, x):
# x shape: (batch, seq_len, num_nodes)
batch_size, seq_len, num_nodes = x.shape
# Reshape to process each sensor independently
x = x.permute(0, 2, 1).reshape(batch_size * num_nodes, seq_len, 1)
# GRU forward
out, _ = self.gru(x)
# Take last timestep output
out = out[:, -1, :] # (batch*nodes, hidden)
out = self.fc(out) # (batch*nodes, 1)
out = out.reshape(batch_size, num_nodes)
return out
What does this code do?
- The permute and reshape in forward() flatten all 207 sensors into a single batch dimension. Each sensor's 12-timestep sequence is processed independently by the same GRU.
- The GRU outputs a hidden state at each timestep, but only the final hidden state (the one that has "seen" the full hour of input) is passed through a linear layer to produce the speed prediction.
- The output is reshaped back to (batch_size, num_nodes) so each sample in the batch gets one prediction per sensor.
Write the training loop
The training function handles batching, loss computation, and validation tracking. It also accepts optional edge_index and edge_weight parameters so you can reuse it for the graph-aware model in the next step.
- Below the GRUBaseline class in the same cell, start adding the train_model function by pasting this code:
# Training function
def train_model(model, X_train, Y_train, X_val, Y_val, epochs=50, lr=0.001,
batch_size=32, edge_index=None, edge_weight=None):
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
criterion = nn.MSELoss()
X_train_t = torch.tensor(X_train, dtype=torch.float32).to(device)
Y_train_t = torch.tensor(Y_train, dtype=torch.float32).to(device)
X_val_t = torch.tensor(X_val, dtype=torch.float32).to(device)
Y_val_t = torch.tensor(Y_val, dtype=torch.float32).to(device)
train_losses, val_losses = [], []
n_batches = len(X_train_t) // batch_size
for epoch in range(epochs):
model.train()
epoch_loss = 0
for i in range(n_batches):
start = i * batch_size
end = start + batch_size
batch_x = X_train_t[start:end]
batch_y = Y_train_t[start:end]
if edge_index is not None:
pred = model(batch_x, edge_index, edge_weight)
else:
pred = model(batch_x)
loss = criterion(pred, batch_y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
epoch_loss += loss.item()
- Continue the train_model function directly below. This section handles validation after each epoch and prints progress every 10 epochs:
avg_train_loss = epoch_loss / n_batches
train_losses.append(avg_train_loss)
# Validation
model.eval()
with torch.no_grad():
if edge_index is not None:
val_pred = model(X_val_t, edge_index, edge_weight)
else:
val_pred = model(X_val_t)
val_loss = criterion(val_pred, Y_val_t).item()
val_losses.append(val_loss)
if (epoch + 1) % 10 == 0:
print(f"Epoch {epoch+1}/{epochs} | Train Loss: {avg_train_loss:.4f} | Val Loss: {val_loss:.4f}")
return train_losses, val_losses
What does the training function do?
- Each epoch iterates over mini-batches of 32 samples. For each batch, the model predicts the next timestep's speeds, and MSELoss measures the squared difference between predictions and actuals.
- The if edge_index is not None check lets the same function train both the GRU baseline (no graph) and the T-GCN (with graph) later.
- After each epoch, the model evaluates on the validation set with gradients disabled (torch.no_grad()) to track whether it is overfitting.
Train and evaluate the baseline
With the model and training function defined, it is time to train the GRU for 50 epochs and see how well a temporal-only model predicts traffic speed.
- Below the train_model function in the same cell, add the code that instantiates, trains, and evaluates the GRU baseline:
# Train GRU baseline
gru_model = GRUBaseline(num_nodes=NUM_NODES, hidden_dim=32).to(device)
print("Training GRU Baseline...")
print(gru_model)
print(f"Parameters: {sum(p.numel() for p in gru_model.parameters()):,}")
gru_train_losses, gru_val_losses = train_model(
gru_model, X_train, Y_train, X_val, Y_val, epochs=50
)
# Evaluate GRU on test set
gru_model.eval()
X_test_t = torch.tensor(X_test, dtype=torch.float32).to(device)
Y_test_t = torch.tensor(Y_test, dtype=torch.float32).to(device)
with torch.no_grad():
gru_preds = gru_model(X_test_t)
gru_preds_denorm = gru_preds.cpu().numpy() * std_speed + mean_speed
Y_test_denorm = Y_test * std_speed + mean_speed
gru_mae = np.mean(np.abs(gru_preds_denorm - Y_test_denorm))
print(f"\nGRU Baseline Test MAE: {gru_mae:.2f} mph")
Before you run this cell, think for a moment: the model has only seen each sensor's own history with no knowledge of neighboring sensors. How accurate do you think it will be at predicting the next 5 minutes of traffic speed?
- Run the cell by pressing Shift+Enter or clicking the play button.
What should I see?
The cell first prints the model architecture (a GRU layer and a Linear layer) followed by the total parameter count. Then training begins.
Every 10 epochs, you see a progress line showing training and validation loss. Training takes a few minutes on a T4 GPU. After 50 epochs, the test MAE prints. A value in the range of 4.0 to 5.0 mph is typical for this baseline.
Seeing a CUDA out of memory error?
- Make sure you selected a T4 GPU runtime in Step 1. Go to Runtime > Change runtime type and confirm a GPU is selected.
- If your Colab session has been running for a while, the GPU memory may be fragmented. Go to Runtime > Restart runtime and re-run all cells from the top.
Still stuck? Help me debug this CUDA memory error when training my GRU baseline on the METR-LA dataset in Google Colab.
The MAE tells you the average prediction error in miles per hour. Now let's see how those predictions look visually against the actual traffic speeds.
- Add the plotting code below the evaluation in the same cell:
# Plot GRU predictions vs actuals
sensor_idx = 50
fig, ax = plt.subplots(figsize=(14, 4))
time_window = slice(0, 200)
ax.plot(Y_test_denorm[time_window, sensor_idx], label='Actual', color='black', linewidth=1.5)
ax.plot(gru_preds_denorm[time_window, sensor_idx], label='GRU Prediction', color='blue',
alpha=0.7, linestyle='--')
ax.set_xlabel('Timestep (5 min intervals)')
ax.set_ylabel('Speed (mph)')
ax.set_title(f'GRU Baseline - Sensor {sensor_idx} | Test MAE: {gru_mae:.2f} mph')
ax.legend()
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()
- Re-run the cell to see the prediction plot.
You should see a chart with the actual speeds in black and the GRU predictions as a blue dashed line for Sensor 50 over 200 timesteps.
Look closely at moments where the actual speed drops or recovers sharply. The GRU's predictions follow the general trend, but they lag behind these sudden changes. This is exactly the gap you will close in the next step.
Why do the predictions lag behind sudden changes?
When a traffic jam starts at one sensor, that slowdown often propagates from a nearby sensor upstream. A GRU looking at only Sensor 50's own history cannot anticipate this. It has to wait until the slowdown appears in Sensor 50's data before reacting.
A model that can also see what neighboring sensors are reporting right now would catch the slowdown earlier. That is the core idea behind the T-GCN you will build next.
Plot not showing up?
- Make sure plt.show() is the last line in the plotting section. Colab inline plots require this call.
- Check that the training and evaluation code ran successfully first. The plot depends on gru_preds_denorm and Y_test_denorm being computed in the evaluation block above.
Need help? My GRU baseline plot is not showing in Google Colab after training.
✔️ Awesome, I've got everything!
Great. Confirm you have saved your notebook. Your GRU baseline and training infrastructure are ready for comparison in the next step.
ⓧ I'd like to double check the full code
Scroll through your Cell 3 and confirm it contains these sections, in order:
- The GRUBaseline class with __init__ (GRU + Linear layers) and forward (reshape, GRU, last timestep, linear, reshape back).
- The train_model function with optimizer/criterion setup, a batch training loop, per-epoch validation, and printing every 10 epochs.
- Model instantiation (GRUBaseline(num_nodes=NUM_NODES, hidden_dim=32)), training call, and test evaluation computing gru_mae.
- A matplotlib plot of predictions vs actuals for sensor_idx = 50 over 200 timesteps.
You now have a concrete number to beat. The GRU baseline captures daily patterns, but it cannot anticipate speed changes that propagate from sensor to sensor across the road network. In the next step, you will build a T-GCN that fuses graph convolutions with the GRU to give the model spatial awareness.
Build the T-GCN Model (GCN + GRU Fusion)
Your GRU baseline from the last step processes each sensor in isolation. It has no way of knowing that Sensor 50's sudden slowdown is about to ripple outward to neighboring sensors on the same highway.
The T-GCN (Temporal Graph Convolutional Network) architecture fixes this by fusing a GCN with a GRU. The GCN reads the road network graph to aggregate speed readings from neighboring sensors. The GRU then processes these spatially-enriched features over time.
In this step, get ready to:
- Define a TGCNCell that fuses spatial encoding (GCN) with temporal processing (GRU).
- Build the full TGCN model that loops the cell over 12 input timesteps.
- Inspect the model architecture and compare parameter counts with the GRU baseline.
Define the T-GCN cell
The TGCNCell is the core building block of the T-GCN architecture. At each timestep, every sensor's speed reading passes through a GCN layer that aggregates information from connected sensors on the road graph. Then a GRU cell takes this spatially-enriched output and updates its temporal memory.
- Create a new code cell below your GRU baseline cell and paste the following TGCNCell class:
class TGCNCell(nn.Module):
def __init__(self, input_dim, hidden_dim):
super(TGCNCell, self).__init__()
self.hidden_dim = hidden_dim
# GCN processes spatial features
self.gcn = GCNConv(input_dim, hidden_dim)
# GRU cell processes temporal dynamics
self.gru_cell = nn.GRUCell(hidden_size=hidden_dim, input_size=hidden_dim)
def forward(self, x, edge_index, edge_weight, h):
# x: (num_nodes, input_dim) - current timestep features
# h: (num_nodes, hidden_dim) - hidden state from previous timestep
# Step 1: GCN captures spatial dependencies
spatial = F.relu(self.gcn(x, edge_index, edge_weight))
# Step 2: GRU captures temporal dynamics
h_new = self.gru_cell(spatial, h)
return h_new
What does this code do?
- GCNConv(input_dim, hidden_dim) is a graph convolution layer from PyTorch Geometric. It takes each sensor's feature and mixes it with features from connected sensors in the road graph.
- nn.GRUCell is a single-step GRU. Unlike nn.GRU (which processes full sequences), GRUCell handles one timestep at a time. This lets us insert the GCN step in between each timestep.
- In forward, sensor features flow through GCN first (spatial), then into the GRU cell (temporal). This two-step fusion is the core T-GCN idea.
Define the full T-GCN model
The TGCN model wraps the cell in a loop over 12 input timesteps. At each step, it calls TGCNCell with the current sensor readings and the road graph structure. The hidden state carries forward across timesteps, accumulating both spatial and temporal context.
- Below the TGCNCell class in the same cell, add the TGCN model class:
class TGCN(nn.Module):
def __init__(self, num_nodes, input_dim=1, hidden_dim=32, output_dim=1):
super(TGCN, self).__init__()
self.num_nodes = num_nodes
self.hidden_dim = hidden_dim
self.tgcn_cell = TGCNCell(input_dim, hidden_dim)
self.fc = nn.Linear(hidden_dim, output_dim)
def forward(self, x, edge_index, edge_weight):
# x: (batch, seq_len, num_nodes)
batch_size, seq_len, num_nodes = x.shape
# Initialize hidden state
h = torch.zeros(batch_size * num_nodes, self.hidden_dim).to(x.device)
for t in range(seq_len):
# Get features for all nodes at timestep t
x_t = x[:, t, :] # (batch, num_nodes)
x_t = x_t.reshape(batch_size * num_nodes, 1) # (batch*nodes, 1)
# Expand edge_index for batched graph
batch_edge_index = self._expand_edge_index(edge_index, batch_size, num_nodes).to(x.device)
batch_edge_weight = edge_weight.repeat(batch_size).to(x.device)
h = self.tgcn_cell(x_t, batch_edge_index, batch_edge_weight, h)
# Output prediction
out = self.fc(h) # (batch*nodes, 1)
out = out.reshape(batch_size, num_nodes)
return out
What does the forward method do?
- The loop iterates over each of the 12 input timesteps. At each step, it reshapes sensor readings into per-node features and passes them through the TGCNCell along with the graph structure.
- _expand_edge_index (which you'll add next) replicates the graph for each sample in a training batch. If your batch has 32 samples and 207 nodes each, the expanded graph has 32 x 207 = 6,624 nodes with the same connection pattern repeated.
- After all 12 timesteps, the final hidden state captures both spatial context (from GCN) and temporal context (from GRU). A linear layer maps this to speed predictions for all 207 sensors simultaneously.
Complete the model and inspect it
The TGCN class still needs its helper method for handling batched graphs. You also need to instantiate the model and print its architecture to compare with the GRU baseline.
- Complete the TGCN class by adding the _expand_edge_index helper method below the forward method (still inside the class, with 4-space indentation). Then add the instantiation code below the class at the top level:
def _expand_edge_index(self, edge_index, batch_size, num_nodes):
"""Replicate graph structure for each sample in the batch."""
edge_indices = []
for i in range(batch_size):
edge_indices.append(edge_index + i * num_nodes)
return torch.cat(edge_indices, dim=1)
# Build model
tgcn_model = TGCN(num_nodes=NUM_NODES, hidden_dim=32).to(device)
print("T-GCN Model Architecture:")
print(tgcn_model)
print(f"Parameters: {sum(p.numel() for p in tgcn_model.parameters()):,}")
# Compare parameter counts
gru_params = sum(p.numel() for p in gru_model.parameters())
tgcn_params = sum(p.numel() for p in tgcn_model.parameters())
print(f"\nGRU Baseline parameters: {gru_params:,}")
print(f"T-GCN parameters: {tgcn_params:,}")
What does _expand_edge_index do?
When training with batches, each sample needs its own copy of the 207-node graph. This helper offsets node indices so sample 0 uses nodes 0-206, sample 1 uses nodes 207-413, and so on. This is a standard pattern in PyTorch Geometric for batched graph processing.
Before you run this cell, do you think the T-GCN will have significantly more parameters than the GRU baseline? The GCN layer processes graph structure, so consider how much extra weight that adds.
- Run the cell with Shift+Enter.
You should see the T-GCN architecture printed, showing a TGCNCell containing a GCNConv layer (input channels = 1, output channels = 32) feeding into a GRUCell (input size = 32, hidden size = 32), plus a Linear output layer. The GCNConv layer adds relatively few parameters compared to the GRU components.
Seeing a NameError or IndentationError?
- If you see NameError: name 'GCNConv' is not defined, make sure you ran Cell 1 first. That cell imports GCNConv from PyTorch Geometric.
- If you see an IndentationError, check that _expand_edge_index is indented with 4 spaces (inside the TGCN class), and the instantiation code starting at tgcn_model = TGCN(...) has no indentation (it sits outside the class).
- If you see NameError: name 'TGCNCell' is not defined, make sure TGCNCell is defined above TGCN in the same cell.
Help me debug my T-GCN model definition.
✔️ Awesome, I've got everything!
Great. Double check that you have the TGCNCell class, the TGCN class, and the model instantiation all in the same cell, and that the output shows both model architectures.
ⓧ I'd like to double check the full code
Your Cell 4 should contain these sections in order. Scroll up through your cell and confirm each one is present:
- TGCNCell class with __init__ (defines self.gcn and self.gru_cell) and forward (GCN then GRU cell).
- TGCN class with __init__ (wraps TGCNCell and nn.Linear), forward (loops over timesteps with hidden state), and _expand_edge_index (replicates graph for batches).
- Model instantiation: tgcn_model = TGCN(num_nodes=NUM_NODES, hidden_dim=32).to(device) followed by print statements for the architecture and parameter comparison.
Your T-GCN model is defined and ready to train. Next up, you'll train it on the same data as the GRU baseline and compare their predictions side by side.
Train the T-GCN and Compare Results
In Step 4, you built a T-GCN architecture that fuses GCN spatial encoding with GRU temporal encoding. The architecture is defined, but it has never seen a single data point.
This step is the payoff. You will train the T-GCN on the exact same data as the GRU baseline and compare their predictions head to head. The question driving this entire project finally gets a concrete answer: does giving the model access to the road network graph actually improve traffic speed predictions?
In this step, get ready to:
- Train the T-GCN model using the same setup as the GRU baseline.
- Evaluate both models and compare their test MAE.
- Visualize the improvement with prediction plots, bar charts, and a per-sensor heatmap.
Train the T-GCN and evaluate both models
The training setup is identical to the GRU baseline: MSE loss, Adam optimizer, learning rate of 0.001, and 50 epochs. The only difference is that the T-GCN also receives the graph's edge connections and weights.
Before you run this, think about what you expect. The GRU baseline scored around 4-5 mph MAE by treating each sensor independently. Do you think spatial awareness will make a noticeable difference?
- Create a new code cell in your notebook and paste this training and evaluation code:
# Prepare graph data for training
ei = edge_index.to(device)
ew = edge_weight.to(device)
print("Training T-GCN...")
tgcn_train_losses, tgcn_val_losses = train_model(
tgcn_model, X_train, Y_train, X_val, Y_val, epochs=50,
edge_index=ei, edge_weight=ew
)
# Evaluate T-GCN on test set
tgcn_model.eval()
with torch.no_grad():
tgcn_preds = tgcn_model(X_test_t, ei, ew)
tgcn_preds_denorm = tgcn_preds.cpu().numpy() * std_speed + mean_speed
tgcn_mae = np.mean(np.abs(tgcn_preds_denorm - Y_test_denorm))
# Print comparison
print(f"\n{'='*40}")
print(f" Model Comparison (Test MAE)")
print(f"{'='*40}")
print(f" GRU Baseline: {gru_mae:.2f} mph")
print(f" T-GCN: {tgcn_mae:.2f} mph")
print(f" Improvement: {gru_mae - tgcn_mae:.2f} mph ({(gru_mae - tgcn_mae)/gru_mae*100:.1f}%)")
print(f"{'='*40}")
What does this code do?
- Moves edge_index and edge_weight to the GPU so the graph structure lives on the same device as the model.
- Calls the same train_model function from Step 3, this time passing graph data via edge_index=ei, edge_weight=ew. When those arguments are not None, the function feeds the graph to the model alongside each batch.
- Switches the model to evaluation mode, runs the test set through, denormalizes predictions back to mph, and prints a side-by-side MAE comparison.
- Run the cell. Training takes a few minutes on the T4 GPU.
- You should see loss values printed every 10 epochs, followed by a comparison table. The T-GCN's test MAE should land around 3.0-4.0 mph, lower than the GRU baseline's 4.0-5.0 mph.
Why is the T-GCN more accurate?
The GRU baseline treats each sensor as if it exists in isolation. When a traffic jam starts at one sensor, the GRU at a neighboring sensor has no idea until the jam shows up in that sensor's own history.
The T-GCN's GCNConv layer aggregates features from connected sensors at every timestep. It can see a speed drop spreading across the network before it reaches a given sensor, giving the GRU cell a head start on predicting the change.
Seeing a CUDA out of memory error?
The T-GCN uses more memory than the GRU because it expands the graph for each batch sample. Try restarting the runtime (**Runtime** > **Restart runtime**) and re-running all cells from the top. If the error persists, reduce the batch size in train_model from 32 to 16.
Still stuck? Help me debug a CUDA out of memory error during T-GCN training.
Plot side-by-side predictions for both models
The MAE numbers confirm the T-GCN is better overall. A prediction plot reveals exactly where it differs from the GRU, especially during sudden speed changes.
- Create a new code cell and paste this plotting code:
# Plot 1: Side-by-side predictions
sensor_idx = 50
fig, axes = plt.subplots(2, 1, figsize=(14, 8), sharex=True)
time_window = slice(0, 200)
axes[0].plot(Y_test_denorm[time_window, sensor_idx], label='Actual', color='black', linewidth=1.5)
axes[0].plot(gru_preds_denorm[time_window, sensor_idx], label='GRU', color='blue', alpha=0.7, linestyle='--')
axes[0].set_ylabel('Speed (mph)')
axes[0].set_title(f'GRU Baseline - Sensor {sensor_idx} | MAE: {gru_mae:.2f} mph')
axes[0].legend()
axes[0].grid(True, alpha=0.3)
axes[1].plot(Y_test_denorm[time_window, sensor_idx], label='Actual', color='black', linewidth=1.5)
axes[1].plot(tgcn_preds_denorm[time_window, sensor_idx], label='T-GCN', color='red', alpha=0.7, linestyle='--')
axes[1].set_xlabel('Timestep (5 min intervals)')
axes[1].set_ylabel('Speed (mph)')
axes[1].set_title(f'T-GCN (GCN+GRU) - Sensor {sensor_idx} | MAE: {tgcn_mae:.2f} mph')
axes[1].legend()
axes[1].grid(True, alpha=0.3)
plt.tight_layout()
plt.show()
What should I see in the prediction plots?
The top panel shows the GRU baseline (blue dashed) against actual speeds (black). The bottom panel shows the T-GCN (red dashed) against the same actuals. Both cover sensor 50 over 200 test timesteps.
Look for moments where speed drops sharply. The GRU predictions tend to lag behind these sudden changes, while the T-GCN tracks them more closely. The T-GCN can "see" congestion building at neighboring sensors through the graph before it reaches this sensor.
- Run the cell. Two stacked line plots appear, each overlaying model predictions on actual speeds for sensor 50.
Both plots look identical?
Make sure you are plotting gru_preds_denorm in the top panel and tgcn_preds_denorm in the bottom panel. If both use the same variable, you will see identical predictions.
Need help? My GRU and T-GCN prediction plots look the same. What went wrong?
Visualize the improvement across all sensors
The prediction plot shows one sensor. Two more visualizations reveal the bigger picture: an overall accuracy bar chart with training loss curves, and a per-sensor heatmap that shows where spatial information helps the most.
- Create a new code cell and paste this bar chart and training loss visualization:
# Plot 2: MAE comparison bar chart
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
# Bar chart
models = ['GRU Baseline', 'T-GCN (GCN+GRU)']
maes = [gru_mae, tgcn_mae]
colors = ['#4285F4', '#EA4335']
axes[0].bar(models, maes, color=colors, width=0.5)
axes[0].set_ylabel('Test MAE (mph)')
axes[0].set_title('Model Comparison')
for i, v in enumerate(maes):
axes[0].text(i, v + 0.05, f'{v:.2f}', ha='center', fontweight='bold')
# Training loss comparison
axes[1].plot(gru_train_losses, label='GRU Train', color='blue', alpha=0.7)
axes[1].plot(tgcn_train_losses, label='T-GCN Train', color='red', alpha=0.7)
axes[1].plot(gru_val_losses, label='GRU Val', color='blue', linestyle='--', alpha=0.5)
axes[1].plot(tgcn_val_losses, label='T-GCN Val', color='red', linestyle='--', alpha=0.5)
axes[1].set_xlabel('Epoch')
axes[1].set_ylabel('MSE Loss')
axes[1].set_title('Training Progress')
axes[1].legend()
axes[1].grid(True, alpha=0.3)
plt.tight_layout()
plt.show()
What do these charts show?
- The bar chart gives a direct visual comparison of test MAE. A shorter bar means more accurate predictions.
- The training loss curves show how each model converged over 50 epochs. The T-GCN's loss typically drops faster and settles lower, confirming that graph structure accelerates learning.
- Run the cell. You should see a bar chart on the left comparing MAE values and training loss curves on the right.
- Create one more code cell and paste this per-sensor heatmap code:
# Per-sensor MAE heatmap
gru_per_sensor_mae = np.mean(np.abs(gru_preds_denorm - Y_test_denorm), axis=0)
tgcn_per_sensor_mae = np.mean(np.abs(tgcn_preds_denorm - Y_test_denorm), axis=0)
improvement = gru_per_sensor_mae - tgcn_per_sensor_mae
fig, ax = plt.subplots(figsize=(14, 3))
im = ax.imshow(improvement.reshape(1, -1), cmap='RdYlGn', aspect='auto')
ax.set_xlabel('Sensor Index')
ax.set_title('Per-Sensor MAE Improvement (GRU MAE - T-GCN MAE)\nGreen = T-GCN is better')
ax.set_yticks([])
plt.colorbar(im, ax=ax, label='MAE Improvement (mph)')
plt.tight_layout()
plt.show()
What does the heatmap reveal?
Each column represents one of the 207 sensors. Green means the T-GCN outperformed the GRU at that sensor. Red means the GRU was better.
Sensors with many graph connections (surrounded by other sensors on the highway) tend to show the strongest green. These are exactly the locations where spatial information matters most, because the T-GCN can draw on neighboring traffic conditions to sharpen its predictions.
- Run the cell. A color-coded heatmap strip appears, showing per-sensor improvement across all 207 sensors.
Heatmap looks all one color?
Check that improvement is computed as gru_per_sensor_mae - tgcn_per_sensor_mae (not the other way around). If the values are all near zero, the colorbar will show a flat strip. This can happen if both models converge to similar performance. Try zooming into the colorbar range to see subtle differences.
Still stuck? Help me understand my per-sensor MAE heatmap results.
✔️ Awesome, I've got everything!
Great. Make sure your notebook is saved to Google Drive.
ⓧ I'd like to double check the full code
Your Cell 5 code (which may span multiple Colab cells) should contain the following sections in order. Scroll up and confirm each one is present:
- Graph data preparation: ei = edge_index.to(device) and ew = edge_weight.to(device).
- T-GCN training via train_model with edge_index=ei, edge_weight=ew.
- T-GCN evaluation: tgcn_model.eval(), denormalization to tgcn_preds_denorm, MAE computation, and a printed comparison table.
- Side-by-side prediction plots for sensor 50 over 200 timesteps.
- Bar chart comparing MAE values and training loss curves for both models.
- Per-sensor MAE improvement heatmap using gru_per_sensor_mae - tgcn_per_sensor_mae.
Secret mission
Multi-Horizon Forecasting
Real traffic systems need predictions much further ahead than just 5 minutes. Extend your T-GCN to predict 15, 30, and 60 minutes into the future, then discover how prediction accuracy degrades at longer horizons and whether the spatial model's advantage holds up.
Clean Up Your Resources
Clean Up Your Resources
This project runs entirely on Google Colab's free tier, so there are no ongoing costs. Decide whether to keep your notebook for future reference or remove it entirely.
Resources you used:
- Google Colab notebook (free, no ongoing cost).
- Downloaded dataset files in the Colab runtime (data/METR-LA.csv and data/adj_mx_METR-LA.pkl, temporary and deleted when the runtime disconnects).
Keep everything running
No action needed. Choose this if you want to keep building on your traffic prediction notebook.
- Save the notebook to your Google Drive by clicking File, then Save a copy in Drive. This preserves all your code and outputs.
- The downloaded dataset files live inside the Colab runtime, which is temporary. When you reopen the notebook later, the first cell will re-download them automatically.
Pause - I'll come back to this later
No action needed. Colab runtimes disconnect automatically after idle time, freeing up the GPU and clearing downloaded data.
- Save the notebook to Google Drive if you have not already (File > Save a copy in Drive).
- Close the Colab browser tab. The runtime shuts down on its own after a period of inactivity.
- When you return, reopen the notebook from Google Drive and re-run Cell 1 to reinstall torch_geometric and re-download the dataset.
Delete - I don't want to use this again
Remove all project resources and start fresh.
- In the Colab menu bar, click Runtime, then click Disconnect and delete runtime. This removes the GPU session and all downloaded files (data/METR-LA.csv and data/adj_mx_METR-LA.pkl).
- If you saved the notebook to Google Drive, open Google Drive, find the notebook file, right-click it, and select Move to trash.
- If you never saved the notebook to Drive, close the browser tab. The unsaved notebook and its runtime data are cleaned up automatically.
That's a Wrap!
That's a Wrap!
Nice work. You just built a spatio-temporal deep learning system that predicts traffic speeds across 207 real highway sensors in Los Angeles, and proved that graph-aware models outperform temporal-only baselines on networked sensor data.
You've learned how to:
- Build a complete deep learning pipeline for traffic speed prediction using the real METR-LA benchmark dataset, from loading and normalizing sensor data to training and evaluating models in Google Colab with a GPU runtime.
- Design a T-GCN model that fuses a Graph Convolutional Network for spatial encoding with a Gated Recurrent Unit for temporal encoding, capturing both road-network relationships and time-series dynamics simultaneously.
- Demonstrate quantitatively that incorporating graph structure improves traffic prediction accuracy over a temporal-only GRU baseline, using side-by-side prediction plots, MAE comparisons, and a per-sensor improvement heatmap.
- Secret Mission: Extended the model to multi-horizon forecasting at 15, 30, and 60 minutes ahead and analyzed how prediction accuracy degrades over longer horizons.
Ready to quiz yourself?