Transformerでグラフデータを扱ってみる

#概要
　Transformerエンコーダを使ってグラフデータを扱います．非自己回帰型のエンコーダでグラフのエッジに合わせた特殊なmaskを行います．

ちなみに[こちら](https://towardsdatascience.com/graph-transformer-generalization-of-transformers-to-graphs-ead2448cff8b)で解説されている[論文](https://arxiv.org/abs/2012.09699)ではもっと複雑に扱っており，SelfAttentionのQKVにエッジ情報を加えたQKVEという入力に拡張しています．

今回考えてみた手法は単純である反面，エッジの特徴を埋め込めない欠点があります．できても単方向or両方向くらいです．

#準備
　今回はGoogle Colabで実行しました，ランタイムはGPUに設定しておきましょう．[こちら](https://gist.github.com/ameya98/b193856171d11d37ada46458f60e73e7)を参考にtorch-geometricをインストールします．

#コード
- Transformerエンコーダレイヤーを用意．
- これは[前回の記事](https://qiita.com/shiba_inu_/items/197f0f48587ed12e591f)からのコピーです．

```python:
import copy
import math
import numpy as np

import torch
torch.manual_seed(41)
import torch.optim as optim
import torch.nn as nn
from torch.nn import functional as F

class SoftmaxAttention(nn.Module):
    def __init__(self, head_dim):
        super().__init__()
        self.head_dim = head_dim

    def forward(self, Q, K, V, mask=None):
        logit = torch.einsum("bhld,bhmd->bhlm",Q,K)/math.sqrt(self.head_dim)

        if mask!=None:
            logit = logit + mask[:,None,:,:]

        attention_weight = F.softmax(logit, dim=-1)
        X = torch.einsum("bhlm,bhmd->bhld",attention_weight,V)

        return X

class SelfAttention(nn.Module):
    def __init__(self, dim, head_dim, num_head):
        super().__init__()

        self.dim = dim
        self.head_dim = head_dim
        self.num_head = num_head

        assert dim == head_dim*num_head,print("ASSERT #dim = head_dim * num_head")

        self.W_q = nn.Linear(self.dim, self.dim)
        self.W_k = nn.Linear(self.dim, self.dim)
        self.W_v = nn.Linear(self.dim, self.dim)

        self.attn = SoftmaxAttention(head_dim)

    def forward(self, X, mask):
        Q = self.split_heads(self.W_q(X))
        K = self.split_heads(self.W_k(X))
        V = self.split_heads(self.W_v(X))

        attn_out = self.attn(Q.float(), K.float(), V.float(), mask.float())
        attn_out = self.combine_heads(attn_out)
        return attn_out

    #[batch,head,len,head_dim]->[batch,len,dim]
    def combine_heads(self, X):
        X = X.transpose(1, 2)
        X = X.reshape(X.size(0), X.size(1), self.num_head * self.head_dim)
        return X

    #[batch,len,dim]->[batch,head,len,head_dim]
    def split_heads(self, X):
        X = X.reshape(X.size(0), X.size(1), self.num_head, self.head_dim)
        X = X.transpose(1, 2)
        return X

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)

        position = torch.arange(max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
        pe = torch.zeros(max_len, 1, d_model)
        pe[:, 0, 0::2] = torch.sin(position * div_term)
        pe[:, 0, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe.transpose(0,1))

    #[batch,len,dim]->[batch,len,dim]
    def forward(self, x):
        x = x + self.pe[:,:x.size(1),:]
        return self.dropout(x)


class TransformerEncoderLayer(nn.Module):
    def __init__(self, attn_layer, ff_layer, norm_layer, drop_layer):
        super().__init__()
        self.attn_layer = attn_layer
        self.ff_layer = ff_layer
        self.norm_layer = norm_layer
        self.drop_layer = drop_layer

    def forward(self, x, mask):
        x = self.drop_layer(self.attn_layer(x,mask)) + x
        x = self.norm_layer(x)

        x = self.drop_layer(self.ff_layer(x)) + x
        x = self.norm_layer(x)

        return x
```
- グラフ用のTransformerです．

・xにノードの特徴量[nodes,dim]を入れてバイアス無しの線形変換で埋め込んでいます．
・maskを使ってエッジ情報を入れます．node_iとnode_jにエッジが無い場合はmask[i,j]=-infに設定します．他は0です．

`Input: x[batch,nodes,dim] adj[batch,nodes,nodes]`

```python:
from torch.nn.modules.container import ModuleList
from torch.nn.modules.normalization import LayerNorm

class TransformerGraph(nn.Module):
    def __init__(self, num_layer, dim, ff_dim, head_dim, num_head, feature_size=None, num_classes=None,drop_p=0.1):
        super().__init__()

        self.node_emb = nn.Linear(feature_size,dim,bias=False)
        self.output = nn.Linear(dim,num_classes)
        
        self.encoders = ModuleList([TransformerEncoderLayer(SelfAttention(dim, head_dim, num_head),\
                            nn.Sequential(nn.Linear(dim,ff_dim),nn.ReLU(),nn.Linear(ff_dim,dim)),\
                                LayerNorm(dim, eps=1e-5),\
                                    nn.Dropout(drop_p)) for i in range(num_layer)])
    
    def forward(self, x, adj):
        mask = torch.zeros(adj.shape).to(x.device).masked_fill_(adj==0,float("-inf"))
        x = self.node_emb(x.float())
        for layer in self.encoders:
            x = layer(x,mask)
        return self.output(x)
```
- データセットを用意

```python:
#edge情報から隣接行列を生成
def edge_adjacency_matrix(edge,node_num):
    print(node_num)
    matrix = torch.zeros(node_num,node_num)
    for i,j in zip(edge[0,:],edge[1,:]):
        matrix[i,j]=1
    return matrix

#不必要なノードのラベルを-1にする
def mask_label(label,mask):
    label = label.masked_fill(mask,-1)
    return label

from torch_geometric.datasets import Planetoid
dataset = Planetoid("./", "Cora", split="public")
data = dataset[0]

nodes_feature = data['x'] #[2708, 1433]
nodes_label = data['y'] #[2708]
adjacency_matrix = edge_adjacency_matrix(data['edge_index'],data.num_nodes) #[2, 10556]

train_mask = data["train_mask"] #[2708]
val_mask = data["val_mask"] #[2708]
test_mask = data["test_mask"] #[2708]

train_label = mask_label(nodes_label,train_mask)
val_label = mask_label(nodes_label,val_mask)
test_label = mask_label(nodes_label,test_mask)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = TransformerGraph(num_layer=3, dim=128, ff_dim=256, head_dim=32, num_head=4, feature_size=1433, num_classes=7).to(device)

import torch.optim as optim
optimizer = optim.Adam(model.parameters(),lr=1e-4)
criterion = nn.CrossEntropyLoss(ignore_index=-1) #train~test_maskで-1にしたlabelを無視
```

- 訓練ループを書きます．

```python:
min_loss = float("inf")

nodes_feature = nodes_feature.unsqueeze(0).to(device)
adjacency_matrix = adjacency_matrix.unsqueeze(0).to(device)

for epoch in range(1000):
    model.train()
    optimizer.zero_grad()
    logits = model(nodes_feature,adjacency_matrix)
    loss = criterion(logits.reshape(-1, logits.size(-1)), train_label.reshape(-1).to(device))
    loss.backward()
    optimizer.step()
    print("train_loss:",loss.item())

    model.eval()
    with torch.no_grad():
        logits = model(nodes_feature,adjacency_matrix)
    loss = criterion(logits.reshape(-1, logits.size(-1)), val_label.reshape(-1).to(device))
    print("val_loss:",loss.item())
    if min_loss > loss.item():
        torch.save(model.state_dict(),"./model.bin")
        print("save!")
```

- テストデータで評価

```python:
model.load_state_dict(torch.load("./model.bin"))
model.eval()

with torch.no_grad():
    logits = model(nodes_feature,adjacency_matrix)
preds = logits.argmax(-1)

accuracy = []
for pred,label,mask in zip(preds.squeeze(),nodes_label,test_mask):
    if mask.item():  #test_maskがFalseのラベルはスキップ
        if pred.item()==label.item():
            accuracy.append(1)
        else:
            accuracy.append(0)
print(sum(accuracy)/len(accuracy))
```

　90%以上出てます，コードが間違ってなければよいのですが…．

#備考
　Transformerは`attention-weight [batch,length,length]`の不要な部分が0になるようにmaskしますが，本モデルは`attention-weight [batch,nodes,nodes]`のエッジの無い組にmaskしています．


しかし，この構造ではエッジの種類（単・共有結合など）に関して埋め込むことが出来ないので注意してください．

エッジのある部分に1次元Embeddingか何かで埋め込むのも一手かもしれません．

#まとめ
　Transformerエンコーダ（非自己回帰==No Causal）でグラフデータを扱いました．

#最後に
　誤っている部分等ございましたら，コメント等で優しく指摘して頂けると嬉しいです．（気付かなければ申し訳ありません）