152 lines
4.9 KiB
Python
Executable File
152 lines
4.9 KiB
Python
Executable File
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
"""
|
|
含残差的双层三明治结构模型
|
|
第一层：时间 -> 空间 -> 时间 -> Conv -> Residual差分 -> 输入第二层
|
|
第二层：时间 -> 空间 -> 时间 -> Conv -> 最终输出
|
|
无效但接近
|
|
"""
|
|
|
|
|
|
class DynamicGraphConstructor(nn.Module):
|
|
def __init__(self, node_num, embed_dim):
|
|
super().__init__()
|
|
self.nodevec1 = nn.Parameter(
|
|
torch.randn(node_num, embed_dim), requires_grad=True
|
|
)
|
|
self.nodevec2 = nn.Parameter(
|
|
torch.randn(node_num, embed_dim), requires_grad=True
|
|
)
|
|
|
|
def forward(self):
|
|
# (N, D) @ (D, N) -> (N, N)
|
|
adj = torch.matmul(self.nodevec1, self.nodevec2.T)
|
|
adj = F.relu(adj)
|
|
adj = F.softmax(adj, dim=-1)
|
|
return adj
|
|
|
|
|
|
class GraphConvBlock(nn.Module):
|
|
def __init__(self, input_dim, output_dim):
|
|
super().__init__()
|
|
self.theta = nn.Linear(input_dim, output_dim)
|
|
self.residual = input_dim == output_dim
|
|
if not self.residual:
|
|
self.res_proj = nn.Linear(input_dim, output_dim)
|
|
|
|
def forward(self, x, adj):
|
|
# x: (B, N, C); adj: (N, N)
|
|
res = x
|
|
x = torch.matmul(adj, x) # 空间卷积
|
|
x = self.theta(x)
|
|
# 残差
|
|
x = x + (res if self.residual else self.res_proj(res))
|
|
return F.relu(x)
|
|
|
|
|
|
class MANBA_Block(nn.Module):
|
|
def __init__(self, input_dim, hidden_dim):
|
|
super().__init__()
|
|
self.attn = nn.MultiheadAttention(
|
|
embed_dim=input_dim, num_heads=4, batch_first=True
|
|
)
|
|
self.ffn = nn.Sequential(
|
|
nn.Linear(input_dim, hidden_dim),
|
|
nn.ReLU(),
|
|
nn.Linear(hidden_dim, input_dim),
|
|
)
|
|
self.norm1 = nn.LayerNorm(input_dim)
|
|
self.norm2 = nn.LayerNorm(input_dim)
|
|
|
|
def forward(self, x):
|
|
# x: (B, N, C) 视 N 维为时间序列长度
|
|
res = x
|
|
x_attn, _ = self.attn(x, x, x)
|
|
x = self.norm1(res + x_attn)
|
|
res2 = x
|
|
x_ffn = self.ffn(x)
|
|
x = self.norm2(res2 + x_ffn)
|
|
return x
|
|
|
|
|
|
class SandwichBlock(nn.Module):
|
|
"""
|
|
时间-空间-时间 三明治结构
|
|
输入/输出： (B, N, hidden_dim)
|
|
"""
|
|
|
|
def __init__(self, num_nodes, embed_dim, hidden_dim):
|
|
super().__init__()
|
|
self.manba1 = MANBA_Block(hidden_dim, hidden_dim * 2)
|
|
self.graph_constructor = DynamicGraphConstructor(num_nodes, embed_dim)
|
|
self.gc = GraphConvBlock(hidden_dim, hidden_dim)
|
|
self.manba2 = MANBA_Block(hidden_dim, hidden_dim * 2)
|
|
|
|
def forward(self, h):
|
|
# h: (B, N, hidden_dim)
|
|
h1 = self.manba1(h)
|
|
adj = self.graph_constructor() # (N, N)
|
|
h2 = self.gc(h1, adj)
|
|
h3 = self.manba2(h2)
|
|
return h3
|
|
|
|
|
|
class EXP(nn.Module):
|
|
def __init__(self, args):
|
|
super().__init__()
|
|
self.horizon = args["horizon"]
|
|
self.output_dim = args["output_dim"]
|
|
self.seq_len = args.get("in_len", 12)
|
|
self.hidden_dim = args.get("hidden_dim", 64)
|
|
self.num_nodes = args["num_nodes"]
|
|
self.embed_dim = args.get("embed_dim", 16)
|
|
|
|
# 输入映射: (batch*N, seq_len) -> hidden_dim
|
|
self.input_proj = nn.Linear(self.seq_len, self.hidden_dim)
|
|
# 两层三明治块
|
|
self.sandwich1 = SandwichBlock(self.num_nodes, self.embed_dim, self.hidden_dim)
|
|
self.sandwich2 = SandwichBlock(self.num_nodes, self.embed_dim, self.hidden_dim)
|
|
# 卷积层用于残差处理
|
|
self.res_conv1 = nn.Conv1d(
|
|
in_channels=self.hidden_dim, out_channels=self.hidden_dim, kernel_size=1
|
|
)
|
|
self.res_conv2 = nn.Conv1d(
|
|
in_channels=self.hidden_dim, out_channels=self.hidden_dim, kernel_size=1
|
|
)
|
|
# 输出映射
|
|
self.out_proj = nn.Linear(self.hidden_dim, self.horizon * self.output_dim)
|
|
|
|
def forward(self, x):
|
|
# x: (B, T, N, D_total)
|
|
x_main = x[..., 0] # (B, T, N)
|
|
B, T, N = x_main.shape
|
|
assert T == self.seq_len
|
|
|
|
# 输入投影 (B, T, N) -> (B*N, T) -> (B, N, hidden_dim)
|
|
x_flat = x_main.permute(0, 2, 1).reshape(B * N, T)
|
|
h0 = self.input_proj(x_flat).view(B, N, self.hidden_dim)
|
|
|
|
# 第一层三明治
|
|
h1_sand = self.sandwich1(h0) # (B, N, hidden_dim)
|
|
# 卷积残差 (节点维度视为长度)
|
|
h1_perm = h1_sand.permute(0, 2, 1) # (B, C, N)
|
|
h1_conv = self.res_conv1(h1_perm)
|
|
h1 = h1_conv.permute(0, 2, 1) # (B, N, hidden_dim)
|
|
|
|
# 计算差分残差作为第二层输入
|
|
h2_input = h1 - h0
|
|
|
|
# 第二层三明治
|
|
h2_sand = self.sandwich2(h2_input)
|
|
# 再次卷积处理
|
|
h2_perm = h2_sand.permute(0, 2, 1)
|
|
h2 = self.res_conv2(h2_perm).permute(0, 2, 1)
|
|
|
|
# 输出映射
|
|
out = self.out_proj(h2) # (B, N, H*D_out)
|
|
out = out.view(B, N, self.horizon, self.output_dim)
|
|
out = out.permute(0, 2, 1, 3) # (B, horizon, N, output_dim)
|
|
return out
|