File size: 1,703 Bytes
8a48f11 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 | import torch
import torch.nn as nn
from torch_geometric.nn import GINEConv, GPSConv
class EventGraphFM(nn.Module):
def __init__(self, n_act: int, n_res: int, n_logs: int, d_model: int = 256,
n_layers: int = 6, n_heads: int = 8, act_dim: int = 64, res_dim: int = 32,
log_dim: int = 16, dropout: float = 0.1, edge_dim: int = 2, num_dim: int = 7):
super().__init__()
self.act_emb = nn.Embedding(n_act, act_dim, padding_idx=0)
self.res_emb = nn.Embedding(n_res, res_dim, padding_idx=0)
self.log_emb = nn.Embedding(n_logs, log_dim)
in_dim = act_dim + res_dim + log_dim + num_dim + 1
self.proj = nn.Linear(in_dim, d_model)
self.layers = nn.ModuleList()
for _ in range(n_layers):
mlp = nn.Sequential(
nn.Linear(d_model, 2 * d_model), nn.ReLU(), nn.Linear(2 * d_model, d_model)
)
conv = GINEConv(mlp, edge_dim=edge_dim)
self.layers.append(GPSConv(d_model, conv, heads=n_heads, dropout=dropout, act="relu"))
self.norm = nn.LayerNorm(d_model)
self.act_out = nn.Linear(d_model, n_act)
self.res_out = nn.Linear(d_model, n_res)
def forward(self, gact, gres, num, log_id, edge_index, edge_attr, batch_vec, mask_flag):
x = torch.cat([
self.act_emb(gact),
self.res_emb(gres),
self.log_emb(log_id),
num,
mask_flag.unsqueeze(-1),
], dim=-1)
h = self.proj(x)
for layer in self.layers:
h = layer(h, edge_index, batch=batch_vec, edge_attr=edge_attr)
h = self.norm(h)
return self.act_out(h), self.res_out(h)
|