Spaces:
Running
Running
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import dgl | |
| import dgl.function as fn | |
| import numpy as np | |
| """ | |
| Graph Transformer Layer | |
| """ | |
| """ | |
| Util functions | |
| """ | |
| def src_dot_dst(src_field, dst_field, out_field): | |
| def func(edges): | |
| return {out_field: (edges.src[src_field] - edges.dst[dst_field])} | |
| return func | |
| def scaled_exp(field, scale_constant): | |
| def func(edges): | |
| # clamp for softmax numerical stability | |
| return {field: torch.exp((edges.data[field] / scale_constant).clamp(-5, 5))} | |
| return func | |
| """ | |
| Single Attention Head | |
| """ | |
| class MultiHeadAttentionLayer(nn.Module): | |
| def __init__(self, in_dim, out_dim, num_heads, use_bias): | |
| super().__init__() | |
| self.out_dim = out_dim | |
| self.num_heads = num_heads | |
| if use_bias: | |
| self.Q = nn.Linear(in_dim, out_dim * num_heads, bias=True) | |
| self.K = nn.Linear(in_dim, out_dim * num_heads, bias=True) | |
| self.V = nn.Linear(in_dim, out_dim * num_heads, bias=True) | |
| else: | |
| self.Q = nn.Linear(in_dim, out_dim * num_heads, bias=False) | |
| self.K = nn.Linear(in_dim, out_dim * num_heads, bias=False) | |
| self.V = nn.Linear(in_dim, out_dim * num_heads, bias=False) | |
| self.M1 = nn.Linear(out_dim, out_dim, bias=False) | |
| self.relu = nn.ReLU() | |
| self.M2 = nn.Linear(out_dim, out_dim, bias=False) | |
| def propagate_attention(self, g): | |
| # Compute attention score | |
| g.apply_edges(src_dot_dst("K_h", "Q_h", "vector")) # , edges) | |
| # if torch.sum(torch.isnan(g.edata["vector"])) > 0: | |
| # print("VECTOR ALREADY NAN HERE") | |
| # 0 / 0 | |
| e_data_m1 = self.M1(g.edata["vector"]) | |
| e_data_m1 = self.relu(e_data_m1) | |
| e_data_m1 = self.M2(e_data_m1) | |
| # print("e_data_m1", e_data_m1[0:2]) | |
| g.edata["vector"] = e_data_m1 | |
| g.apply_edges(scaled_exp("vector", np.sqrt(self.out_dim))) | |
| # print("vector", g.edata["vector"][0:2]) | |
| # if torch.sum(torch.isnan(g.edata["vector"])) > 0: | |
| # print(g.edata["vector"]) | |
| # Send weighted values to target nodes | |
| eids = g.edges() | |
| # vector attention to modulate individual channels | |
| g.send_and_recv(eids, fn.u_mul_e("V_h", "vector", "V_h"), fn.sum("V_h", "wV")) | |
| # print("wV", g.ndata["wV"][0:2]) | |
| g.send_and_recv(eids, fn.copy_e("vector", "vector"), fn.sum("vector", "z")) | |
| # print("z", g.ndata["z"][0:2]) | |
| # if torch.sum(torch.isnan(g.ndata["z"])) > 0: | |
| # 0 / 0 | |
| def forward(self, g, h): | |
| Q_h = self.Q(h) | |
| K_h = self.K(h) | |
| V_h = self.V(h) | |
| # if torch.sum(torch.isnan(Q_h)) > 0: | |
| # print("Q_h ALREADY NAN HERE") | |
| # 0 / 0 | |
| # if torch.sum(torch.isnan(V_h)) > 0: | |
| # print("V_h ALREADY NAN HERE") | |
| # 0 / 0 | |
| # if torch.sum(torch.isnan(K_h)) > 0: | |
| # print("K_h ALREADY NAN HERE") | |
| # 0 / 0 | |
| # Reshaping into [num_nodes, num_heads, feat_dim] to | |
| # get projections for multi-head attention | |
| g.ndata["Q_h"] = Q_h.view(-1, self.num_heads, self.out_dim) | |
| g.ndata["K_h"] = K_h.view(-1, self.num_heads, self.out_dim) | |
| g.ndata["V_h"] = V_h.view(-1, self.num_heads, self.out_dim) | |
| # print("q_h", Q_h[0:2]) | |
| # print("K_h", K_h[0:2]) | |
| # print("V_h", V_h[0:2]) | |
| self.propagate_attention(g) | |
| # g.ndata["z"] = g.ndata["z"].tile((1, 1, self.out_dim)) | |
| mask_empty = g.ndata["z"] > 0 | |
| head_out = g.ndata["wV"] | |
| head_out[mask_empty] = head_out[mask_empty] / (g.ndata["z"][mask_empty]) | |
| # g.ndata["z"] = g.ndata["z"][:, :, 0].view( | |
| # g.ndata["wV"].shape[0], self.num_heads, 1 | |
| # ) | |
| # print("head_out", head_out[0:2]) | |
| # if torch.sum(torch.isnan(head_out)) > 0: | |
| # print("head_out ALREADY NAN HERE") | |
| # 0 / 0 | |
| return head_out | |
| class GraphTransformerLayer(nn.Module): | |
| """ | |
| Param: | |
| """ | |
| def __init__( | |
| self, | |
| in_dim, | |
| out_dim, | |
| num_heads, | |
| dropout=0.0, | |
| layer_norm=False, | |
| batch_norm=True, | |
| residual=False, | |
| use_bias=False, | |
| ): | |
| super().__init__() | |
| self.in_channels = in_dim | |
| self.out_channels = out_dim | |
| self.num_heads = num_heads | |
| self.dropout = dropout | |
| self.residual = residual | |
| self.layer_norm = layer_norm | |
| self.batch_norm = batch_norm | |
| self.attention = MultiHeadAttentionLayer( | |
| in_dim, out_dim // num_heads, num_heads, use_bias | |
| ) | |
| self.O = nn.Linear(out_dim, out_dim) | |
| if self.layer_norm: | |
| self.layer_norm1 = nn.LayerNorm(out_dim) | |
| if self.batch_norm: | |
| self.batch_norm1 = nn.BatchNorm1d(out_dim) | |
| # FFN | |
| self.FFN_layer1 = nn.Linear(out_dim, out_dim * 2) | |
| self.FFN_layer2 = nn.Linear(out_dim * 2, out_dim) | |
| if self.layer_norm: | |
| self.layer_norm2 = nn.LayerNorm(out_dim) | |
| if self.batch_norm: | |
| self.batch_norm2 = nn.BatchNorm1d(out_dim) | |
| def forward(self, g, h): | |
| h_in1 = h # for first residual connection | |
| # multi-head attention out | |
| attn_out = self.attention(g, h) | |
| h = attn_out.view(-1, self.out_channels) | |
| # print("output of the attention ", h[0:2]) | |
| # if torch.sum(torch.isnan(h)) > 0: | |
| # print("output of the attention ALREADY NAN HERE") | |
| # 0 / 0 | |
| h = F.dropout(h, self.dropout, training=self.training) | |
| h = self.O(h) | |
| if self.residual: | |
| h = h_in1 + h # residual connection | |
| # print("output of residual ", h[0:2]) | |
| # if torch.sum(torch.isnan(h)) > 0: | |
| # print("output of the residual ALREADY NAN HERE") | |
| # 0 / 0 | |
| if self.layer_norm: | |
| h = self.layer_norm1(h) | |
| if self.batch_norm: | |
| h = self.batch_norm1(h) | |
| # # print("output of batchnorm ", h[0:2]) | |
| # if torch.sum(torch.isnan(h)) > 0: | |
| # print("output of the batchnorm ALREADY NAN HERE") | |
| # 0 / 0 | |
| h_in2 = h # for second residual connection | |
| # FFN | |
| h = self.FFN_layer1(h) | |
| h = F.relu(h) | |
| h = F.dropout(h, self.dropout, training=self.training) | |
| h = self.FFN_layer2(h) | |
| # print("output of FFN_layer2 ", h[0:2]) | |
| # if torch.sum(torch.isnan(h)) > 0: | |
| # print("output of the FFN_layer2 ALREADY NAN HERE") | |
| # 0 / 0 | |
| if self.residual: | |
| h = h_in2 + h # residual connection | |
| if self.layer_norm: | |
| h = self.layer_norm2(h) | |
| if self.batch_norm: | |
| h = self.batch_norm2(h) | |
| return h | |
| def __repr__(self): | |
| return "{}(in_channels={}, out_channels={}, heads={}, residual={})".format( | |
| self.__class__.__name__, | |
| self.in_channels, | |
| self.out_channels, | |
| self.num_heads, | |
| self.residual, | |
| ) | |