Spaces:
Running
Running
File size: 2,284 Bytes
d015578 |
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 |
from copy import deepcopy
import torch
from torch import nn
import torch.nn.functional as F
from spiga.models.gnn.layers import MLP
class GAT(nn.Module):
def __init__(self, input_dim: int, output_dim: int, num_heads=4):
super().__init__()
num_heads_in = num_heads
self.reshape = None
if input_dim != output_dim:
for num_heads_in in range(num_heads, 0, -1):
if input_dim % num_heads_in == 0:
break
self.reshape = MLP([input_dim, output_dim])
self.attention = MessagePassing(input_dim, num_heads_in, out_dim=output_dim)
def forward(self, features):
message, prob = self.attention(features)
if self.reshape:
features = self.reshape(features)
output = features + message
return output, prob
class MessagePassing(nn.Module):
def __init__(self, feature_dim: int, num_heads: int, out_dim=None):
super().__init__()
self.attn = Attention(num_heads, feature_dim)
self.mlp = MLP([feature_dim*2, feature_dim*2, out_dim])
def forward(self, features):
message, prob = self.attn(features, features, features)
return self.mlp(torch.cat([features, message], dim=1)), prob
class Attention(nn.Module):
def __init__(self, num_heads: int, feature_dim: int):
super().__init__()
assert feature_dim % num_heads == 0
self.dim = feature_dim // num_heads
self.num_heads = num_heads
self.merge = nn.Conv1d(feature_dim, feature_dim, kernel_size=1)
self.proj = nn.ModuleList([deepcopy(self.merge) for _ in range(3)])
def forward(self, query, key, value):
batch_dim = query.size(0)
query, key, value = [l(x).view(batch_dim, self.dim, self.num_heads, -1)
for l, x in zip(self.proj, (query, key, value))]
x, prob = self.attention(query, key, value)
return self.merge(x.contiguous().view(batch_dim, self.dim*self.num_heads, -1)), prob
def attention(self, query, key, value):
dim = query.shape[1]
scores = torch.einsum('bdhn,bdhm->bhnm', query, key) / dim ** .5
prob = F.softmax(scores, dim=-1)
return torch.einsum('bhnm,bdhm->bdhn', prob, value), prob
|