308 lines
14 KiB
Python
308 lines
14 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
from torch.nn import init
|
||
from data.get_adj import get_adj
|
||
import numbers
|
||
|
||
|
||
# --- 基础算子 ---
|
||
class NConv(nn.Module):
|
||
def forward(self, x, adj):
|
||
return torch.einsum('ncwl,vw->ncvl', (x, adj)).contiguous()
|
||
|
||
|
||
class DyNconv(nn.Module):
|
||
def forward(self, x, adj):
|
||
return torch.einsum('ncvl,nvwl->ncwl', (x, adj)).contiguous()
|
||
|
||
|
||
class Linear(nn.Module):
|
||
def __init__(self, c_in, c_out, bias=True):
|
||
super().__init__()
|
||
self.mlp = nn.Conv2d(c_in, c_out, kernel_size=1, bias=bias)
|
||
|
||
def forward(self, x):
|
||
return self.mlp(x)
|
||
|
||
|
||
class Prop(nn.Module):
|
||
def __init__(self, c_in, c_out, gdep, dropout, alpha):
|
||
super().__init__()
|
||
self.nconv = NConv()
|
||
self.mlp = Linear(c_in, c_out)
|
||
self.gdep, self.dropout, self.alpha = gdep, dropout, alpha
|
||
|
||
def forward(self, x, adj):
|
||
adj = adj + torch.eye(adj.size(0), device=x.device)
|
||
d = adj.sum(1)
|
||
a = adj / d.view(-1, 1)
|
||
h = x
|
||
for _ in range(self.gdep):
|
||
h = self.alpha * x + (1 - self.alpha) * self.nconv(h, a)
|
||
return self.mlp(h)
|
||
|
||
|
||
class MixProp(nn.Module):
|
||
def __init__(self, c_in, c_out, gdep, dropout, alpha):
|
||
super().__init__()
|
||
self.nconv = NConv()
|
||
self.mlp = Linear((gdep + 1) * c_in, c_out)
|
||
self.gdep, self.dropout, self.alpha = gdep, dropout, alpha
|
||
|
||
def forward(self, x, adj):
|
||
adj = adj + torch.eye(adj.size(0), device=x.device)
|
||
d = adj.sum(1)
|
||
a = adj / d.view(-1, 1)
|
||
out = [x]
|
||
h = x
|
||
for _ in range(self.gdep):
|
||
h = self.alpha * x + (1 - self.alpha) * self.nconv(h, a)
|
||
out.append(h)
|
||
return self.mlp(torch.cat(out, dim=1))
|
||
|
||
|
||
class DyMixprop(nn.Module):
|
||
def __init__(self, c_in, c_out, gdep, dropout, alpha):
|
||
super().__init__()
|
||
self.nconv = DyNconv()
|
||
self.mlp1 = Linear((gdep + 1) * c_in, c_out)
|
||
self.mlp2 = Linear((gdep + 1) * c_in, c_out)
|
||
self.gdep, self.dropout, self.alpha = gdep, dropout, alpha
|
||
self.lin1, self.lin2 = Linear(c_in, c_in), Linear(c_in, c_in)
|
||
|
||
def forward(self, x):
|
||
x1 = torch.tanh(self.lin1(x))
|
||
x2 = torch.tanh(self.lin2(x))
|
||
adj = self.nconv(x1.transpose(2, 1), x2)
|
||
adj0 = torch.softmax(adj, dim=2)
|
||
adj1 = torch.softmax(adj.transpose(2, 1), dim=2)
|
||
# 两条分支
|
||
out1, out2 = [x], [x]
|
||
h = x
|
||
for _ in range(self.gdep):
|
||
h = self.alpha * x + (1 - self.alpha) * self.nconv(h, adj0)
|
||
out1.append(h)
|
||
h = x
|
||
for _ in range(self.gdep):
|
||
h = self.alpha * x + (1 - self.alpha) * self.nconv(h, adj1)
|
||
out2.append(h)
|
||
return self.mlp1(torch.cat(out1, dim=1)) + self.mlp2(torch.cat(out2, dim=1))
|
||
|
||
|
||
class DilatedInception(nn.Module):
|
||
def __init__(self, cin, cout, dilation_factor=2):
|
||
super().__init__()
|
||
self.kernels = [2, 3, 6, 7]
|
||
cout_each = int(cout / len(self.kernels))
|
||
self.convs = nn.ModuleList([nn.Conv2d(cin, cout_each, kernel_size=(1, k), dilation=(1, dilation_factor))
|
||
for k in self.kernels])
|
||
|
||
def forward(self, x):
|
||
outs = [conv(x)[..., -self.convs[-1](x).size(3):] for conv in self.convs]
|
||
return torch.cat(outs, dim=1)
|
||
|
||
|
||
class GraphConstructor(nn.Module):
|
||
def __init__(self, nnodes, k, dim, device, alpha=3, static_feat=None):
|
||
super().__init__()
|
||
self.nnodes, self.k, self.dim, self.alpha, self.device = nnodes, k, dim, alpha, device
|
||
self.static_feat = static_feat
|
||
if static_feat is not None:
|
||
xd = static_feat.shape[1]
|
||
self.lin1, self.lin2 = nn.Linear(xd, dim), nn.Linear(xd, dim)
|
||
else:
|
||
self.emb1 = nn.Embedding(nnodes, dim)
|
||
self.emb2 = nn.Embedding(nnodes, dim)
|
||
self.lin1, self.lin2 = nn.Linear(dim, dim), nn.Linear(dim, dim)
|
||
|
||
def forward(self, idx):
|
||
if self.static_feat is None:
|
||
vec1, vec2 = self.emb1(idx), self.emb2(idx)
|
||
else:
|
||
vec1 = vec2 = self.static_feat[idx, :]
|
||
vec1 = torch.tanh(self.alpha * self.lin1(vec1))
|
||
vec2 = torch.tanh(self.alpha * self.lin2(vec2))
|
||
a = torch.mm(vec1, vec2.transpose(1, 0)) - torch.mm(vec2, vec1.transpose(1, 0))
|
||
adj = F.relu(torch.tanh(self.alpha * a))
|
||
mask = torch.zeros(idx.size(0), idx.size(0), device=self.device)
|
||
s1, t1 = adj.topk(self.k, 1)
|
||
mask.scatter_(1, t1, s1.new_ones(s1.size()))
|
||
return adj * mask
|
||
|
||
|
||
class LayerNorm(nn.Module):
|
||
__constants__ = ['normalized_shape', 'eps', 'elementwise_affine']
|
||
|
||
def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True):
|
||
super().__init__()
|
||
if isinstance(normalized_shape, numbers.Integral):
|
||
normalized_shape = (normalized_shape,)
|
||
self.normalized_shape, self.eps, self.elementwise_affine = tuple(normalized_shape), eps, elementwise_affine
|
||
if elementwise_affine:
|
||
self.weight = nn.Parameter(torch.Tensor(*normalized_shape))
|
||
self.bias = nn.Parameter(torch.Tensor(*normalized_shape))
|
||
init.ones_(self.weight);
|
||
init.zeros_(self.bias)
|
||
else:
|
||
self.register_parameter('weight', None)
|
||
self.register_parameter('bias', None)
|
||
|
||
def forward(self, x, idx):
|
||
if self.elementwise_affine:
|
||
return F.layer_norm(x, tuple(x.shape[1:]), self.weight[:, idx, :], self.bias[:, idx, :], self.eps)
|
||
else:
|
||
return F.layer_norm(x, tuple(x.shape[1:]), self.weight, self.bias, self.eps)
|
||
|
||
def extra_repr(self):
|
||
return f'{self.normalized_shape}, eps={self.eps}, elementwise_affine={self.elementwise_affine}'
|
||
|
||
|
||
# --- 合并后的模型类，支持 teacher 与 stmlp 两种分支 ---
|
||
class STMLP(nn.Module):
|
||
def __init__(self, args):
|
||
super().__init__()
|
||
# 参数从字典中读取
|
||
self.adj_mx = get_adj(args)
|
||
self.num_nodes = args['num_nodes']
|
||
self.feature_dim = args['input_dim']
|
||
|
||
self.input_window = args['input_window']
|
||
self.output_window = args['output_window']
|
||
self.output_dim = args['output_dim']
|
||
self.device = args['device']
|
||
|
||
self.gcn_true = args['gcn_true']
|
||
self.buildA_true = args['buildA_true']
|
||
self.gcn_depth = args['gcn_depth']
|
||
self.dropout = args['dropout']
|
||
self.subgraph_size = args['subgraph_size']
|
||
self.node_dim = args['node_dim']
|
||
self.dilation_exponential = args['dilation_exponential']
|
||
|
||
self.conv_channels = args['conv_channels']
|
||
self.residual_channels = args['residual_channels']
|
||
self.skip_channels = args['skip_channels']
|
||
self.end_channels = args['end_channels']
|
||
|
||
self.layers = args['layers']
|
||
self.propalpha = args['propalpha']
|
||
self.tanhalpha = args['tanhalpha']
|
||
self.layer_norm_affline = args['layer_norm_affline']
|
||
|
||
self.model_type = args['model_type'] # 'teacher' 或 'stmlp'
|
||
self.idx = torch.arange(self.num_nodes).to(self.device)
|
||
self.predefined_A = None if self.adj_mx is None else (torch.tensor(self.adj_mx) - torch.eye(self.num_nodes)).to(
|
||
self.device)
|
||
self.static_feat = None
|
||
|
||
# transformer（保留原有结构）
|
||
self.encoder_layer = nn.TransformerEncoderLayer(d_model=12, nhead=4, batch_first=True)
|
||
self.transformer_encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=3)
|
||
|
||
# 构建各层
|
||
self.start_conv = nn.Conv2d(self.feature_dim, self.residual_channels, kernel_size=1)
|
||
self.gc = GraphConstructor(self.num_nodes, self.subgraph_size, self.node_dim, self.device, alpha=self.tanhalpha,
|
||
static_feat=self.static_feat)
|
||
# 计算 receptive_field
|
||
kernel_size = 7
|
||
if self.dilation_exponential > 1:
|
||
self.receptive_field = int(
|
||
self.output_dim + (kernel_size - 1) * (self.dilation_exponential ** self.layers - 1) / (
|
||
self.dilation_exponential - 1))
|
||
else:
|
||
self.receptive_field = self.layers * (kernel_size - 1) + self.output_dim
|
||
|
||
self.filter_convs = nn.ModuleList()
|
||
self.gate_convs = nn.ModuleList()
|
||
self.residual_convs = nn.ModuleList()
|
||
self.skip_convs = nn.ModuleList()
|
||
self.norm = nn.ModuleList()
|
||
self.stu_mlp = nn.ModuleList([nn.Sequential(nn.Linear(c, c), nn.Linear(c, c), nn.Linear(c, c))
|
||
for c in [13, 7, 1]])
|
||
if self.gcn_true:
|
||
self.gconv1 = nn.ModuleList()
|
||
self.gconv2 = nn.ModuleList()
|
||
|
||
new_dilation = 1
|
||
for i in range(1):
|
||
rf_size_i = int(1 + i * (kernel_size - 1) * (self.dilation_exponential ** self.layers - 1) / (
|
||
self.dilation_exponential - 1)) if self.dilation_exponential > 1 else i * self.layers * (
|
||
kernel_size - 1) + 1
|
||
for j in range(1, self.layers + 1):
|
||
rf_size_j = int(rf_size_i + (kernel_size - 1) * (self.dilation_exponential ** j - 1) / (
|
||
self.dilation_exponential - 1)) if self.dilation_exponential > 1 else rf_size_i + j * (
|
||
kernel_size - 1)
|
||
self.filter_convs.append(
|
||
DilatedInception(self.residual_channels, self.conv_channels, dilation_factor=new_dilation))
|
||
self.gate_convs.append(
|
||
DilatedInception(self.residual_channels, self.conv_channels, dilation_factor=new_dilation))
|
||
self.residual_convs.append(nn.Conv2d(self.conv_channels, self.residual_channels, kernel_size=1))
|
||
k_size = (1, self.input_window - rf_size_j + 1) if self.input_window > self.receptive_field else (
|
||
1, self.receptive_field - rf_size_j + 1)
|
||
self.skip_convs.append(nn.Conv2d(self.conv_channels, self.skip_channels, kernel_size=k_size))
|
||
if self.gcn_true:
|
||
self.gconv1.append(MixProp(self.conv_channels, self.residual_channels, self.gcn_depth, self.dropout,
|
||
self.propalpha))
|
||
self.gconv2.append(MixProp(self.conv_channels, self.residual_channels, self.gcn_depth, self.dropout,
|
||
self.propalpha))
|
||
norm_size = (self.residual_channels, self.num_nodes,
|
||
self.input_window - rf_size_j + 1) if self.input_window > self.receptive_field else (
|
||
self.residual_channels, self.num_nodes, self.receptive_field - rf_size_j + 1)
|
||
self.norm.append(LayerNorm(norm_size, elementwise_affine=self.layer_norm_affline))
|
||
new_dilation *= self.dilation_exponential
|
||
|
||
self.end_conv_1 = nn.Conv2d(self.skip_channels, self.end_channels, kernel_size=1, bias=True)
|
||
self.end_conv_2 = nn.Conv2d(self.end_channels, self.output_window, kernel_size=1, bias=True)
|
||
k0 = (1, self.input_window) if self.input_window > self.receptive_field else (1, self.receptive_field)
|
||
self.skip0 = nn.Conv2d(self.feature_dim, self.skip_channels, kernel_size=k0, bias=True)
|
||
kE = (1, self.input_window - self.receptive_field + 1) if self.input_window > self.receptive_field else (1, 1)
|
||
self.skipE = nn.Conv2d(self.residual_channels, self.skip_channels, kernel_size=kE, bias=True)
|
||
# 最后输出分支，根据模型类型选择不同的头
|
||
if self.model_type == 'teacher':
|
||
self.tt_linear1 = nn.Linear(self.residual_channels, self.input_window)
|
||
self.tt_linear2 = nn.Linear(1, 32)
|
||
self.ss_linear1 = nn.Linear(self.residual_channels, self.input_window)
|
||
self.ss_linear2 = nn.Linear(1, 32)
|
||
else: # stmlp
|
||
self.out_linear1 = nn.Linear(self.residual_channels, self.input_window)
|
||
self.out_linear2 = nn.Linear(1, 32)
|
||
|
||
def forward(self, source, idx=None):
|
||
source = source[..., 0:1]
|
||
sout, tout = [], []
|
||
inputs = source.transpose(1, 3)
|
||
assert inputs.size(3) == self.input_window, 'input sequence length mismatch'
|
||
if self.input_window < self.receptive_field:
|
||
inputs = F.pad(inputs, (self.receptive_field - self.input_window, 0, 0, 0))
|
||
if self.gcn_true:
|
||
adp = self.gc(self.idx if idx is None else idx) if self.buildA_true else self.predefined_A
|
||
x = self.start_conv(inputs)
|
||
skip = self.skip0(F.dropout(inputs, self.dropout, training=self.training))
|
||
for i in range(self.layers):
|
||
residual = x
|
||
filters = torch.tanh(self.filter_convs[i](x))
|
||
gate = torch.sigmoid(self.gate_convs[i](x))
|
||
x = F.dropout(filters * gate, self.dropout, training=self.training)
|
||
tout.append(x)
|
||
s = self.skip_convs[i](x)
|
||
skip = s + skip
|
||
if self.gcn_true:
|
||
x = self.gconv1[i](x, adp) + self.gconv2[i](x, adp.transpose(1, 0))
|
||
else:
|
||
x = self.stu_mlp[i](x)
|
||
x = x + residual[:, :, :, -x.size(3):]
|
||
x = self.norm[i](x, self.idx if idx is None else idx)
|
||
sout.append(x)
|
||
skip = self.skipE(x) + skip
|
||
x = F.relu(skip)
|
||
x = F.relu(self.end_conv_1(x))
|
||
x = self.end_conv_2(x)
|
||
if self.model_type == 'teacher':
|
||
ttout = self.tt_linear2(self.tt_linear1(tout[-1].transpose(1, 3)).transpose(1, 3))
|
||
ssout = self.ss_linear2(self.ss_linear1(sout[-1].transpose(1, 3)).transpose(1, 3))
|
||
return x, ttout, ssout
|
||
else:
|
||
x_ = self.out_linear2(self.out_linear1(tout[-1].transpose(1, 3)).transpose(1, 3))
|
||
return x, x_, x
|