File size: 4,800 Bytes
2ab60a1 |
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 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 |
import torch
from torch import nn
from torch_geometric.nn import MessagePassing
class CustomTransformer(nn.Module):
def __init__(self, feat_dim, nhead, num_encoder_layers, dim_feedforward, dropout, first_seq=1, second_seq=1):
super(CustomTransformer, self).__init__()
self.seq_len = first_seq + second_seq + 2
self.first_seq = first_seq
self.second_seq = second_seq
encoder_layer = nn.TransformerEncoderLayer(d_model=feat_dim, nhead=nhead,
dim_feedforward=dim_feedforward,
dropout=dropout, batch_first=True)
self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_encoder_layers,
enable_nested_tensor=False)
self.cls_token_param = nn.Parameter(torch.ones(1, 1, feat_dim))
self.sep_token_param = nn.Parameter(torch.zeros(1, 1, feat_dim))
self.pos_param = nn.Parameter(torch.zeros(1, self.seq_len, feat_dim))
def forward(self, *x):
# x1: (Be (!), H) | x1: (Be, H)
# x2: (Be, H) | x2: (Be, H)
# x3 : -- | x3: (Be, H)
first_seq = [x.unsqueeze(1) for x in x[:self.first_seq]] # (Be, 1, H)
second_seq = [x.unsqueeze(1) for x in x[self.first_seq:]] # (Be, 1, H)
cls_token = self.cls_token_param.expand(first_seq[0].size(0), -1, -1)
sep_token = self.sep_token_param.expand(first_seq[0].size(0), -1, -1)
x = torch.cat([cls_token] + first_seq + [sep_token] + second_seq, dim=1)
x += self.pos_param
x = self.transformer_encoder(x)
return x[:, 0, :]
class GNNTransformModel(MessagePassing):
def __init__(self, num_node_features, num_edge_features, dropout_rate=.1, hid_dim=128):
super(GNNTransformModel, self).__init__(aggr='add')
self.node_encoder = nn.Sequential(
nn.Linear(num_node_features, hid_dim),
nn.Linear(hid_dim, hid_dim),
nn.LayerNorm(hid_dim),
)
self.edge_encoder = nn.Sequential(
nn.Linear(num_edge_features, hid_dim),
nn.Linear(hid_dim, hid_dim),
nn.LayerNorm(hid_dim),
)
self.node_decoder = nn.Linear(hid_dim, num_node_features)
self.edge_decoder = nn.Linear(hid_dim, num_edge_features)
self.node_message_passing = CustomTransformer(
hid_dim, 4, 4, hid_dim, dropout_rate,
first_seq=1
)
self.edge_message_passing = CustomTransformer(
hid_dim, 4, 4, hid_dim, dropout_rate,
first_seq=2
)
# self.node_message_passing = nn.Sequential(
# nn.Linear(hid_dim * 2, hid_dim),
# nn.LeakyReLU(),
# nn.Dropout(p=dropout_rate),
# nn.Linear(hid_dim, hid_dim),
# nn.LeakyReLU(),
# nn.Dropout(p=dropout_rate),
# nn.Linear(hid_dim, hid_dim),
# nn.LeakyReLU(),
# )
# self.edge_message_passing = nn.Sequential(
# nn.Linear(hid_dim * 3, hid_dim),
# nn.LeakyReLU(),
# nn.Dropout(p=dropout_rate),
# nn.Linear(hid_dim, hid_dim),
# nn.LeakyReLU(),
# nn.Dropout(p=dropout_rate),
# nn.Linear(hid_dim, hid_dim),
# nn.LeakyReLU(),
# )
def forward(self, x, edge_index, edge_attr):
# x: (Bn, N) Bn ~ 4978
# edge_attr: (Be, E) Be ~ 10708
# edge_index: (2, Be)
x = self.node_encoder(x) # (Bn, H)
edge_attr = self.edge_encoder(edge_attr) if len(edge_attr) > 0 else edge_attr # (Be, H)
# Handling graphs with no edges
if len(edge_index) > 0 and edge_index.shape[1] > 0:
x = self.propagate(edge_index, x=x, edge_attr=edge_attr)
edge_attr = self.edge_updater(edge_index, x=x, edge_attr=edge_attr) if len(edge_attr) > 0 else edge_attr
out_node_features = self.node_decoder(x)
out_edge_features = self.edge_decoder(edge_attr) if len(edge_attr) > 0 else edge_attr
return out_node_features, out_edge_features
def message(self, x_j, edge_attr):
# x_i: --
# x_j: (Be (!), H)
# edge_attr: (Be, H)
return self.node_message_passing(x_j, edge_attr)
def edge_update(self, x_i, x_j, edge_attr):
# x_i: (Be, H)
# x_j: (Be, H)
# edge_attr: (Be, H)
return self.edge_message_passing(x_i, x_j, edge_attr)
def update(self, aggr_out):
# Only for nodes
# aggr_out: (Bn, H)
return aggr_out
|