czzhangheng
/
TrafficWheel
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
TrafficWheel/model/STEP/discrete_graph_learning.py

184 lines
8.8 KiB
Python
Raw Blame History

# Discrete Graph Learning
import torch
import numpy as np
from torch import nn
import torch.nn.functional as F
import os
from .similarity import batch_cosine_similarity, batch_dot_similarity
def sample_gumbel(shape, eps=1e-20, device=None):
uniform = torch.rand(shape).to(device)
return -torch.autograd.Variable(torch.log(-torch.log(uniform + eps) + eps))
def gumbel_softmax_sample(logits, temperature, eps=1e-10):
sample = sample_gumbel(logits.size(), eps=eps, device=logits.device)
y = logits + sample
return F.softmax(y / temperature, dim=-1)
def gumbel_softmax(logits, temperature, hard=False, eps=1e-10):
"""Sample from the Gumbel-Softmax distribution and optionally discretize.
Args:
logits: [batch_size, n_class] unnormalized log-probs
temperature: non-negative scalar
hard: if True, take argmax, but differentiate w.r.t. soft sample y
Returns:
[batch_size, n_class] sample from the Gumbel-Softmax distribution.
If hard=True, then the returned sample will be one-hot, otherwise it will
be a probabilitiy distribution that sums to 1 across classes
"""
y_soft = gumbel_softmax_sample(logits, temperature=temperature, eps=eps)
if hard:
shape = logits.size()
_, k = y_soft.data.max(-1)
y_hard = torch.zeros(*shape).to(logits.device)
y_hard = y_hard.zero_().scatter_(-1, k.view(shape[:-1] + (1,)), 1.0)
y = torch.autograd.Variable(y_hard - y_soft.data) + y_soft
else:
y = y_soft
return y
class DiscreteGraphLearning(nn.Module):
"""Dynamic graph learning module."""
def __init__(self, dataset_name, k, input_seq_len, output_seq_len):
super().__init__()
self.k = k # the "k" of knn graph
self.num_nodes = {"METR-LA": 207, "PEMS04": 307, "PEMS03": 358, "PEMS-BAY": 325, "PEMS07": 883, "PEMS08": 170}[dataset_name]
self.train_length = {"METR-LA": 23990, "PEMS04": 13599, "PEMS03": 15303, "PEMS07": 16513, "PEMS-BAY": 36482, "PEMS08": 14284}[dataset_name]
# 尝试加载数据，如果文件不存在则使用默认值
try:
data_path = f"data/{dataset_name}/data_in{input_seq_len}_out{output_seq_len}.pkl"
if os.path.exists(data_path):
import pickle
with open(data_path, 'rb') as f:
data = pickle.load(f)
self.node_feats = torch.from_numpy(data["processed_data"]).float()[:self.train_length, :, 0]
else:
# 如果文件不存在，创建一个随机数据作为占位符
print(f"Warning: Data file {data_path} not found. Using random data as placeholder.")
self.node_feats = torch.randn(self.train_length, self.num_nodes, 1)
except Exception as e:
print(f"Warning: Failed to load data for {dataset_name}. Using random data as placeholder. Error: {e}")
self.node_feats = torch.randn(self.train_length, self.num_nodes, 1)
# CNN for global feature extraction
## for the dimension, see https://github.com/zezhishao/STEP/issues/1#issuecomment-1191640023
self.dim_fc = {"METR-LA": 383552, "PEMS04": 217296, "PEMS03": 244560, "PEMS07": 263920, "PEMS-BAY": 583424, "PEMS08": 228256}[dataset_name]
self.embedding_dim = 100
## network structure
self.conv1 = torch.nn.Conv1d(1, 8, 10, stride=1) # .to(device)
self.conv2 = torch.nn.Conv1d(8, 16, 10, stride=1) # .to(device)
self.fc = torch.nn.Linear(self.dim_fc, self.embedding_dim)
self.bn1 = torch.nn.BatchNorm1d(8)
self.bn2 = torch.nn.BatchNorm1d(16)
self.bn3 = torch.nn.BatchNorm1d(self.embedding_dim)
# FC for transforming the features from TSFormer
## for the dimension, see https://github.com/zezhishao/STEP/issues/1#issuecomment-1191640023
self.dim_fc_mean = {"METR-LA": 16128, "PEMS-BAY": 16128, "PEMS03": 16128 * 2, "PEMS04": 16128 * 2, "PEMS07": 16128, "PEMS08": 16128 * 2}[dataset_name]
self.fc_mean = nn.Linear(self.dim_fc_mean, 100)
# discrete graph learning
self.fc_cat = nn.Linear(self.embedding_dim, 2)
self.fc_out = nn.Linear((self.embedding_dim) * 2, self.embedding_dim)
self.dropout = nn.Dropout(0.5)
def encode_one_hot(labels):
# reference code https://github.com/chaoshangcs/GTS/blob/8ed45ff1476639f78c382ff09ecca8e60523e7ce/model/pytorch/model.py#L149
classes = set(labels)
classes_dict = {c: np.identity(len(classes))[i, :] for i, c in enumerate(classes)}
labels_one_hot = np.array(list(map(classes_dict.get, labels)), dtype=np.int32)
return labels_one_hot
self.rel_rec = torch.FloatTensor(np.array(encode_one_hot(np.where(np.ones((self.num_nodes, self.num_nodes)))[0]), dtype=np.float32))
self.rel_send = torch.FloatTensor(np.array(encode_one_hot(np.where(np.ones((self.num_nodes, self.num_nodes)))[1]), dtype=np.float32))
def get_k_nn_neighbor(self, data, k=11*207, metric="cosine"):
"""
data: tensor B, N, D
metric: cosine or dot
"""
if metric == "cosine":
batch_sim = batch_cosine_similarity(data, data)
elif metric == "dot":
batch_sim = batch_dot_similarity(data, data) # B, N, N
else:
assert False, "unknown metric"
batch_size, num_nodes, _ = batch_sim.shape
adj = batch_sim.view(batch_size, num_nodes*num_nodes)
res = torch.zeros_like(adj)
top_k, indices = torch.topk(adj, k, dim=-1)
res.scatter_(-1, indices, top_k)
adj = torch.where(res != 0, 1.0, 0.0).detach().clone()
adj = adj.view(batch_size, num_nodes, num_nodes)
adj.requires_grad = False
return adj
def forward(self, long_term_history, tsformer):
"""Learning discrete graph structure based on TSFormer.
Args:
long_term_history (torch.Tensor): very long-term historical MTS with shape [B, P * L, N, C], which is used in the TSFormer.
P is the number of segments (patches), and L is the length of segments (patches).
tsformer (nn.Module): the pre-trained TSFormer.
Returns:
torch.Tensor: Bernoulli parameter (unnormalized) of each edge of the learned dependency graph. Shape: [B, N * N, 2].
torch.Tensor: the output of TSFormer with shape [B, N, P, d].
torch.Tensor: the kNN graph with shape [B, N, N], which is used to guide the training of the dependency graph.
torch.Tensor: the sampled graph with shape [B, N, N].
"""
device = long_term_history.device
batch_size, _, num_nodes, _ = long_term_history.shape
# generate global feature
global_feat = self.node_feats.to(device).transpose(1, 0).view(num_nodes, 1, -1)
global_feat = self.bn2(F.relu(self.conv2(self.bn1(F.relu(self.conv1(global_feat))))))
global_feat = global_feat.view(num_nodes, -1)
global_feat = F.relu(self.fc(global_feat))
global_feat = self.bn3(global_feat)
global_feat = global_feat.unsqueeze(0).expand(batch_size, num_nodes, -1) # Gi in Eq. (2)
# generate dynamic feature based on TSFormer
hidden_states = tsformer(long_term_history[..., [0]])
# The dynamic feature has now been removed,
# as we found that it could lead to instability in the learning of the underlying graph structure.
# dynamic_feat = F.relu(self.fc_mean(hidden_states.reshape(batch_size, num_nodes, -1))) # relu(FC(Hi)) in Eq. (2)
# time series feature
node_feat = global_feat
# learning discrete graph structure
receivers = torch.matmul(self.rel_rec.to(node_feat.device), node_feat)
senders = torch.matmul(self.rel_send.to(node_feat.device), node_feat)
edge_feat = torch.cat([senders, receivers], dim=-1)
edge_feat = torch.relu(self.fc_out(edge_feat))
# Bernoulli parameter (unnormalized) Theta_{ij} in Eq. (2)
bernoulli_unnorm = self.fc_cat(edge_feat)
# sampling
## differentiable sampling via Gumbel-Softmax in Eq. (4)
sampled_adj = gumbel_softmax(bernoulli_unnorm, temperature=0.5, hard=True)
sampled_adj = sampled_adj[..., 0].clone().reshape(batch_size, num_nodes, -1)
## remove self-loop
mask = torch.eye(num_nodes, num_nodes).unsqueeze(0).bool().to(sampled_adj.device)
sampled_adj.masked_fill_(mask, 0)
# prior graph based on TSFormer
adj_knn = self.get_k_nn_neighbor(hidden_states.reshape(batch_size, num_nodes, -1), k=self.k*self.num_nodes, metric="cosine")
mask = torch.eye(num_nodes, num_nodes).unsqueeze(0).bool().to(adj_knn.device)
adj_knn.masked_fill_(mask, 0)
return bernoulli_unnorm, hidden_states, adj_knn, sampled_adj
Reference in New Issue View Git Blame Copy Permalink