122 lines
4.0 KiB
Python
122 lines
4.0 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class DynamicGraphConstructor(nn.Module):
|
|
def __init__(self, node_num, embed_dim):
|
|
super().__init__()
|
|
self.nodevec1 = nn.Parameter(torch.randn(node_num, embed_dim))
|
|
self.nodevec2 = nn.Parameter(torch.randn(node_num, embed_dim))
|
|
|
|
def forward(self):
|
|
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):
|
|
res = x
|
|
x = torch.matmul(adj, x)
|
|
x = self.theta(x)
|
|
if self.residual:
|
|
x = x + res
|
|
else:
|
|
x = x + 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):
|
|
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 EXPExpert(nn.Module): # 原 EXP 改名
|
|
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.graph = DynamicGraphConstructor(self.num_nodes, embed_dim=16)
|
|
self.input_proj = nn.Linear(self.seq_len, self.hidden_dim)
|
|
self.gc = GraphConvBlock(self.hidden_dim, self.hidden_dim)
|
|
self.manba = MANBA_Block(self.hidden_dim, self.hidden_dim * 2)
|
|
self.out_proj = nn.Linear(self.hidden_dim, self.horizon * self.output_dim)
|
|
|
|
def forward(self, x):
|
|
x = x[..., 0] # (B, T, N)
|
|
B, T, N = x.shape
|
|
x = x.permute(0, 2, 1).reshape(B * N, T)
|
|
h = self.input_proj(x).view(B, N, -1)
|
|
adj = self.graph()
|
|
h = self.gc(h, adj)
|
|
h = self.manba(h)
|
|
out = self.out_proj(h)
|
|
return out.view(B, N, self.horizon, self.output_dim).permute(0, 2, 1, 3)
|
|
|
|
|
|
class EXP(nn.Module):
|
|
def __init__(self, args, num_experts=4, top_k=2):
|
|
super().__init__()
|
|
self.num_experts = num_experts
|
|
self.top_k = top_k
|
|
self.experts = nn.ModuleList([EXPExpert(args) for _ in range(num_experts)])
|
|
self.gate = nn.Sequential(
|
|
nn.Linear(args['in_len'] * args['num_nodes'], 128),
|
|
nn.ReLU(),
|
|
nn.Linear(128, num_experts)
|
|
)
|
|
|
|
def forward(self, x):
|
|
B = x.size(0)
|
|
# Flatten input for gating
|
|
gate_input = x[..., 0].reshape(B, -1) # (B, T*N)
|
|
gate_logits = self.gate(gate_input) # (B, num_experts)
|
|
gate_scores = F.softmax(gate_logits, dim=-1) # soft selection
|
|
|
|
# Get top-k experts
|
|
topk_val, topk_idx = torch.topk(gate_scores, self.top_k, dim=-1) # (B, k)
|
|
|
|
outputs = torch.zeros_like(self.experts[0](x)) # (B, H, N, D_out)
|
|
for i in range(self.top_k):
|
|
idx = topk_idx[:, i]
|
|
for expert_id in torch.unique(idx):
|
|
mask = idx == expert_id
|
|
if mask.sum() == 0:
|
|
continue
|
|
selected_x = x[mask]
|
|
expert_output = self.experts[expert_id](selected_x)
|
|
outputs[mask] += topk_val[mask, i].unsqueeze(1).unsqueeze(1).unsqueeze(1) * expert_output
|
|
|
|
return outputs # (B, H, N, D_out)
|