深入探讨时空图神经网络的原理与Pytorch实现，展示数据转化与模型预测能力。
原文标题：时空图神经网络原理及Pytorch实现
原文作者：数据派THU
冷月清谈：
怜星夜思：
2、除了股市，ST-GNN还可以应用在哪些领域？
3、ST-GNN的实现是否会面临什么具体挑战？
原文内容
来源：算法进阶本文约4700字，建议阅读8分钟
本文我们将深入探讨这些模型的原理，并实现一个相对简单的示例，以更深入地理解它们的能力和应用。
图神经网络（GNN）
时空图神经网络 （ST-GNN）
ST-GNN的Pytorch实现
import yfinance as yf import datetime as dt import pandas as pd from sklearn.preprocessing import StandardScalerimport plotly.graph_objs as go
from plotly.offline import iplot
import matplotlib.pyplot as plt
############ Dataset download #################
start_date = dt.datetime(2013,1,1)
end_date = dt.datetime(2024,3,7)
#loading from yahoo finance
google = yf.download(“GOOGL”,start_date, end_date)
apple = yf.download(“AAPL”,start_date, end_date)
Microsoft = yf.download(“MSFT”, start_date, end_date)
Amazon = yf.download(“AMZN”, start_date, end_date)
meta = yf.download(“META”, start_date, end_date)
Nvidia = yf.download(“NVDA”, start_date, end_date)
data = pd.DataFrame({‘google’: google[‘Open’],‘microsoft’: Microsoft[‘Open’],‘amazon’: Amazon[‘Open’],
‘Nvidia’: Nvidia[‘Open’],‘meta’: meta[‘Open’], ‘apple’: apple[‘Open’]})
############## Scaling data ######################
scaler = StandardScaler()
data_scaled = pd.DataFrame(scaler.fit_transform(data), columns=data.columns)
-
邻接矩阵的定义：AdjacencyMatrix 函数定义了图的邻接矩阵（连通性），这通常是基于手头物理系统的结构来完成的。然而，在这里，作者仅使用了一个全1矩阵，即所有节点都与所有其他节点相连。
-
股市数据集类：StockMarketDataset 类旨在为训练时空图神经网络（ST-GNNs）创建数据集。这个类中包含的方法有：
-
数据序列生成：DatasetCreate 方法生成数据序列。
-
构造图边：_create_edges 方法使用邻接矩阵构造图的边。
-
生成数据序列：_create_sequences 方法通过在输入的股市数据上滑动窗口来生成数据序列。
def AdjacencyMatrix(L): AdjM = np.ones((L,L)) return AdjM
class StockMarketDataset:
def init(self, W,N_hist, N_pred):
self.W = W
self.N_hist = N_hist
self.N_pred = N_pred
def DatasetCreate(self):
num_days, self.n_node = data_scaled.shape
n_window = self.N_hist + self.N_pred
edge_index, edge_attr = self._create_edges(self.n_node)
sequences = self._create_sequences(data_scaled, self.n_node, n_window, edge_index, edge_attr)
return sequences
def _create_edges(self, n_node):
edge_index = torch.zeros((2, n_node2), dtype=torch.long)
edge_attr = torch.zeros((n_node2, 1))
num_edges = 0
for i in range(n_node):
for j in range(n_node):
if self.W[i, j] != 0:
edge_index[:, num_edges] = torch.tensor([i, j], dtype=torch.long)
edge_attr[num_edges, 0] = self.W[i, j]
num_edges += 1
edge_index = edge_index[:, :num_edges]
edge_attr = edge_attr[:num_edges]
return edge_index, edge_attr
def _create_sequences(self, data, n_node, n_window, edge_index, edge_attr):
sequences =
num_days, _ = data.shape
for i in range(num_days):
sta = i
end = i+n_window
full_window = np.swapaxes(data[sta:end, :], 0, 1)
g = Data(x=torch.FloatTensor(full_window[:, :self.N_hist]),
y=torch.FloatTensor(full_window[:, self.N_hist:]),
edge_index=edge_index,
num_nodes=n_node)
sequences.append(g)
return sequences
from torch_geometric.loader import DataLoaderdef train_val_test_splits(sequences, splits):
total = len(sequences)
split_train, split_val, split_test = splitsCalculate split indices
idx_train = int(total * split_train)
idx_val = int(total * (split_train + split_val))
indices = [i for i in range(len(sequences)-100)]
random.shuffle(indices)
train = [sequences[index] for index in indices[:idx_train]]
val = [sequences[index] for index in indices[idx_train:idx_val]]
test = [sequences[index] for index in indices[idx_val:]]
return train, val, test
‘’‘Setting up the hyper paramaters’‘’
n_nodes = 6
n_hist = 50
n_pred = 10
batch_size = 32Adjacency matrix
W = AdjacencyMatrix(n_nodes)
transorm data into graphical time series
dataset = StockMarketDataset(W, n_hist, n_pred)
sequences = dataset.DatasetCreate()train, validation, test split
splits = (0.9, 0.05, 0.05)
train, val, test = train_val_test_splits(sequences, splits)
train_dataloader = DataLoader(train, batch_size=batch_size, shuffle=True, drop_last = True)
val_dataloader = DataLoader(val, batch_size=batch_size, shuffle=True, drop_last=True)
test_dataloader = DataLoader(test, batch_size=batch_size, shuffle=True, drop_last = True)
import torch import torch.nn.functional as F from torch_geometric.nn import GATConvclass ST_GNN_Model(torch.nn.Module):
def init(self, in_channels, out_channels, n_nodes,gru_hs_l1, gru_hs_l2, heads=1, dropout=0.01):
super(ST_GAT, self).init()
self.n_pred = out_channels
self.heads = heads
self.dropout = dropout
self.n_nodes = n_nodes
self.gru_hidden_size_l1 = gru_hs_l1
self.gru_hidden_size_l2 = gru_hs_l2
self.decoder_hidden_size = self.gru_hidden_size_l2enconder GRU layers
self.gat = GATConv(in_channels=in_channels, out_channels=in_channels,
heads=heads, dropout=dropout, concat=False)
self.encoder_gru_l1 = torch.nn.GRU(input_size=self.n_nodes,
hidden_size=self.gru_hidden_size_l1, num_layers=1,
bias = True)
self.encoder_gru_l2 = torch.nn.GRU(input_size=self.gru_hidden_size_l1,
hidden_size=self.gru_hidden_size_l2, num_layers = 1,
bias = True)
self.GRU_decoder = torch.nn.GRU(input_size = self.gru_hidden_size_l2, hidden_size = self.decoder_hidden_size,
num_layers =1, bias = True, dropout= self.dropout)self.prediction_layer = torch.nn.Linear(self.decoder_hidden_size, self.n_nodes*self.n_pred, bias= True)
def forward(self, data, device):
x, edge_index = data.x, data.edge_index
if device == ‘cpu’:
x = torch.FloatTensor(x)
else:
x = torch.cuda.FloatTensor(x)
x = self.gat(x, edge_index)
x = F.dropout(x, self.dropout, training=self.training)
batch_size = data.num_graphs
n_node = int(data.num_nodes / batch_size)
x = torch.reshape(x, (batch_size, n_node, data.num_features))
x = torch.movedim(x, 2, 0)
encoderl1_outputs, _ = self.encoder_gru_l1(x)
x = F.relu(encoderl1_outputs)
encoderl2_outputs, h2 = self.encoder_gru_l2(x)
x = F.relu(encoderl2_outputs)
x, _ = self.GRU_decoder(x,h2)
x = torch.squeeze(x[-1,:,:])
x = self.prediction_layer(x)
x = torch.reshape(x, (batch_size, self.n_nodes, self.n_pred))
x = torch.reshape(x, (batch_size*self.n_nodes, self.n_pred))
return x
import torch import torch.optim as optimHyperparameters
gru_hs_l1 = 16
gru_hs_l2 = 16
learning_rate = 1e-3
Epochs = 50
device = ‘cuda’ if torch.cuda.is_available() else ‘cpu’
model = ST_GNN_Model(in_channels=n_hist, out_channels=n_pred, n_nodes=n_nodes, gru_hs_l1=gru_hs_l1, gru_hs_l2 = gru_hs_l2)
pretrained = False
model_path = “ST_GNN_Model.pth”
if pretrained:
model.load_state_dict(torch.load(model_path))
optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=1e-7)
criterion = torch.nn.MSELoss()
model.to(device)
for epoch in range(Epochs):
model.train()
for _, batch in enumerate(tqdm(train_dataloader, desc=f"Epoch {epoch}“)):
batch = batch.to(device)
optimizer.zero_grad()
y_pred = torch.squeeze(model(batch, device))
loss= criterion(y_pred.float(), torch.squeeze(batch.y).float())
loss.backward()
optimizer.step()
print(f"Loss: {loss:.7f}”)
@torch.no_grad()
def Extract_results(model, device, dataloader, type=‘’):
model.eval()
model.to(device)
n = 0Evaluate model on all data
for i, batch in enumerate(dataloader):
batch = batch.to(device)
if batch.x.shape[0] == 1:
pass
else:
with torch.no_grad():
pred = model(batch, device)
truth = batch.y.view(pred.shape)
if i == 0:
y_pred = torch.zeros(len(dataloader), pred.shape[0], pred.shape[1])
y_truth = torch.zeros(len(dataloader), pred.shape[0], pred.shape[1])
y_pred[i, :pred.shape[0], :] = pred
y_truth[i, :pred.shape[0], :] = truth
n += 1
y_pred_flat = torch.reshape(y_pred, (len(dataloader),batch_size,n_nodes,n_pred))
y_truth_flat = torch.reshape(y_truth,(len(dataloader),batch_size,n_nodes,n_pred))
return y_pred_flat, y_truth_flat
def plot_results(predictions,actual, step, node):
predictions = torch.tensor(predictions[:,:,node,step]).squeeze()
actual = torch.tensor(actual[:,:,node,step]).squeeze()
pred_values_float = torch.reshape(predictions,(-1,))
actual_values_float = torch.reshape(actual, (-1,))
scatter_trace = go.Scatter(
x=actual_values_float,
y=pred_values_float,
mode=‘markers’,
marker=dict(
size=10,
opacity=0.5,
color=‘rgba(255,255,255,0)’,
line=dict(
width=2,
color=‘rgba(152, 0, 0, .8)’,
)
),
name=‘Actual vs Predicted’
)
line_trace = go.Scatter(
x=[min(actual_values_float), max(actual_values_float)],
y=[min(actual_values_float), max(actual_values_float)],
mode=‘lines’,
marker=dict(color=‘blue’),
name=‘Perfect Prediction’
)
data = [scatter_trace, line_trace]
layout = dict(
title=‘Actual vs Predicted Values’,
xaxis=dict(title=‘Actual Values’),
yaxis=dict(title=‘Predicted Values’),
autosize=False,
width=800,
height=600
)
fig = dict(data=data, layout=layout)
iplot(fig)
y_pred, y_truth = Extract_results(model, device, test_dataloader, ‘Test’)
plot_results(y_pred, y_truth,9,0) # timestep, node
@torch.no_grad() def forecastModel(model, device, dataloader, node): model.eval() model.to(device) for i, batch in enumerate(dataloader): batch = batch.to(device) with torch.no_grad(): pred = model(batch, device) truth = batch.y.view(pred.shape) # the shape should [batch_size, nodes, number of predictions] truth = torch.reshape(truth, [batch_size, n_nodes,n_pred]) pred = torch.reshape(pred, [batch_size, n_nodes,n_pred]) x = batch.x x = torch.reshape(x, [batch_size, n_nodes,n_hist])y_pred = torch.squeeze(pred[0, node, :])
y_truth = torch.squeeze(truth[0,node,:])
y_past = torch.squeeze(x[0, node, :])
t_range = [t for t in range(len(y_past))]
break
t_shifted = [t_range[-1]+1+t for t in range(len(y_pred))]
trace1 = go.Scatter(x =t_range, y= y_past, mode = “markers”, name = “Historical data”)
trace2 = go.Scatter(x=t_shifted, y=y_pred, mode = “markers”, name = “pred”)
trace3 = go.Scatter(x=t_shifted, y=y_truth, mode = “markers”, name = “truth”)
layout = go.Layout(title = “forecasting”, xaxis=dict(title = ‘time’),
yaxis=dict(title = ‘y-value’), width = 1000, height = 600)
figure = go.Figure(data = [trace1, trace2, trace3], layout = layout)
iplot(figure)
forecastModel(model, device, test_dataloader, 0)
总结
编辑：王菁
