| 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) |
|
|
