PyG (PyTorch Geometric) で GAT (Graph Attention Networks)

グラフ構造を深層学習する PyG (PyTorch Geometric)  を Google Colaboratory 上で使ってみました。まずは GAT (Graph Attention Networks) を用いて、node property prediction （頂点のラベルを予測）です。

# PyG (PyTorch Geometric) インストール

PyG (PyTorch Geometric) のレポジトリは https://github.com/pyg-team/pytorch_geometric にあります。


```python
import os
import torch
os.environ['TORCH'] = torch.__version__
print(torch.__version__)

!pip install -q torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}.html
!pip install -q torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}.html
!pip install -q git+https://github.com/pyg-team/pytorch_geometric.git

import torch_geometric
```

    1.12.0+cu113
    [K     |████████████████████████████████| 7.9 MB 2.9 MB/s 
    [K     |████████████████████████████████| 3.5 MB 2.7 MB/s 
    [?25h  Building wheel for torch-geometric (setup.py) ... [?25l[?25hdone


# データセットのロード

ベンチマークなどで使えるデータセットは https://pytorch-geometric.readthedocs.io/en/latest/modules/datasets.html に置いてあります。


```python
from torch_geometric import datasets
```

## KarateClub データセット

その中で最もデータサイズの小さい KarateClub データセットを試しに使ってみましょう。


```python
dataset = datasets.KarateClub()
dataset
```




    KarateClub()



次のようにして、データの内容を外観できます。


```python
print("number of graphs:\t\t",len(dataset))
print("number of classes:\t\t",dataset.num_classes)
print("number of node features:\t",dataset.num_node_features)
print("number of edge features:\t",dataset.num_edge_features)
```

    number of graphs:		 1
    number of classes:		 4
    number of node features:	 34
    number of edge features:	 0


データの本体は .data にあります。次のようにして、中に入っているデータの形状（shape）を外観できます。


```python
dataset.data
```




    Data(x=[34, 34], edge_index=[2, 156], y=[34], train_mask=[34])



詳細は次の通りです。


```python
print("edge_index:\t\t",dataset.data.edge_index.shape)
print(dataset.data.edge_index)
print("\n")
print("train_mask:\t\t",dataset.data.train_mask.shape)
print(dataset.data.train_mask)
print("\n")
print("x:\t\t",dataset.data.x.shape)
print(dataset.data.x)
print("\n")
print("y:\t\t",dataset.data.y.shape)
print(dataset.data.y)
```

    edge_index:		 torch.Size([2, 156])
    tensor([[ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  1,  1,
              1,  1,  1,  1,  1,  1,  1,  2,  2,  2,  2,  2,  2,  2,  2,  2,  2,  3,
              3,  3,  3,  3,  3,  4,  4,  4,  5,  5,  5,  5,  6,  6,  6,  6,  7,  7,
              7,  7,  8,  8,  8,  8,  8,  9,  9, 10, 10, 10, 11, 12, 12, 13, 13, 13,
             13, 13, 14, 14, 15, 15, 16, 16, 17, 17, 18, 18, 19, 19, 19, 20, 20, 21,
             21, 22, 22, 23, 23, 23, 23, 23, 24, 24, 24, 25, 25, 25, 26, 26, 27, 27,
             27, 27, 28, 28, 28, 29, 29, 29, 29, 30, 30, 30, 30, 31, 31, 31, 31, 31,
             31, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 33, 33, 33, 33, 33,
             33, 33, 33, 33, 33, 33, 33, 33, 33, 33, 33, 33],
            [ 1,  2,  3,  4,  5,  6,  7,  8, 10, 11, 12, 13, 17, 19, 21, 31,  0,  2,
              3,  7, 13, 17, 19, 21, 30,  0,  1,  3,  7,  8,  9, 13, 27, 28, 32,  0,
              1,  2,  7, 12, 13,  0,  6, 10,  0,  6, 10, 16,  0,  4,  5, 16,  0,  1,
              2,  3,  0,  2, 30, 32, 33,  2, 33,  0,  4,  5,  0,  0,  3,  0,  1,  2,
              3, 33, 32, 33, 32, 33,  5,  6,  0,  1, 32, 33,  0,  1, 33, 32, 33,  0,
              1, 32, 33, 25, 27, 29, 32, 33, 25, 27, 31, 23, 24, 31, 29, 33,  2, 23,
             24, 33,  2, 31, 33, 23, 26, 32, 33,  1,  8, 32, 33,  0, 24, 25, 28, 32,
             33,  2,  8, 14, 15, 18, 20, 22, 23, 29, 30, 31, 33,  8,  9, 13, 14, 15,
             18, 19, 20, 22, 23, 26, 27, 28, 29, 30, 31, 32]])
    
    
    train_mask:		 torch.Size([34])
    tensor([ True, False, False, False,  True, False, False, False,  True, False,
            False, False, False, False, False, False, False, False, False, False,
            False, False, False, False,  True, False, False, False, False, False,
            False, False, False, False])
    
    
    x:		 torch.Size([34, 34])
    tensor([[1., 0., 0.,  ..., 0., 0., 0.],
            [0., 1., 0.,  ..., 0., 0., 0.],
            [0., 0., 1.,  ..., 0., 0., 0.],
            ...,
            [0., 0., 0.,  ..., 1., 0., 0.],
            [0., 0., 0.,  ..., 0., 1., 0.],
            [0., 0., 0.,  ..., 0., 0., 1.]])
    
    
    y:		 torch.Size([34])
    tensor([1, 1, 1, 1, 3, 3, 3, 1, 0, 1, 3, 1, 1, 1, 0, 0, 3, 1, 0, 1, 0, 1, 0, 0,
            2, 2, 0, 0, 2, 0, 0, 2, 0, 0])


NetworkX を使って可視化します。ノード（頂点）は、目的変数 y に応じて色分けしています。


```python
import matplotlib.pyplot as plt
import networkx as nx
from torch_geometric.utils.convert import to_networkx, from_networkx

G = to_networkx(dataset.data, to_undirected=False)
color = ["r", "g", "m", "c", "y"]

plt.figure(figsize=(8, 8))
pos = nx.spring_layout(G, k=0.3)
node_color = [color[y] for y in dataset.data.y]
nx.draw_networkx_nodes(G, pos, alpha=0.5, node_color=node_color)
nx.draw_networkx_edges(G, pos, alpha=0.1)
nx.draw_networkx_labels(G, pos)
plt.axis("off")
plt.savefig("G.png")
plt.show()
```

![Qiita_PyG_(PyTorch_Geometric)_で_GAT_(Graph_Attention_Networks)_15_0.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/211162/dca10a52-3c67-74f5-2239-79fc07b0188d.png)

    

    


## PubMed データセット

次の例として、生命科学や生物医学に関する文献データベースである PubMed データセットを見てみます。


```python
dataset = datasets.Planetoid(root=".", name="PubMed")
dataset
```

    Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.pubmed.x
    Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.pubmed.tx
    Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.pubmed.allx
    Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.pubmed.y
    Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.pubmed.ty
    Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.pubmed.ally
    Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.pubmed.graph
    Downloading https://github.com/kimiyoung/planetoid/raw/master/data/ind.pubmed.test.index
    Processing...
    Done!





    PubMed()



次のようにして、データの内容を外観できます。


```python
print("number of graphs:\t\t",len(dataset))
print("number of classes:\t\t",dataset.num_classes)
print("number of node features:\t",dataset.num_node_features)
print("number of edge features:\t",dataset.num_edge_features)
```

    number of graphs:		 1
    number of classes:		 3
    number of node features:	 500
    number of edge features:	 0


データの本体は .data にあります。次のようにして、中に入っているデータの形状（shape）を外観できます。


```python
dataset.data
```




    Data(x=[19717, 500], edge_index=[2, 88648], y=[19717], train_mask=[19717], val_mask=[19717], test_mask=[19717])



詳細は次の通りです。


```python
print("edge_index:\t\t",dataset.data.edge_index.shape)
print(dataset.data.edge_index)
print("\n")
print("x:\t\t",dataset.data.x.shape)
print(dataset.data.x)
print("\n")
print("y:\t\t",dataset.data.y.shape)
print(dataset.data.y)
```

    edge_index:		 torch.Size([2, 88648])
    tensor([[    0,     0,     0,  ..., 19714, 19715, 19716],
            [ 1378,  1544,  6092,  ..., 12278,  4284, 16030]])
    
    
    x:		 torch.Size([19717, 500])
    tensor([[0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
            [0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
            [0.1046, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
            ...,
            [0.0000, 0.0194, 0.0080,  ..., 0.0000, 0.0000, 0.0000],
            [0.1078, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
            [0.0000, 0.0266, 0.0000,  ..., 0.0000, 0.0000, 0.0000]])
    
    
    y:		 torch.Size([19717])
    tensor([1, 1, 0,  ..., 2, 0, 2])


NetworkX でネットワークを描画したいところですが、描画してもよく分からないので省略します。

# train_val_test_split

ノードラベルを train, val, test に分割する train_val_test_split を次のように実装します。


```python
def train_val_test_split(data, val_ratio: float = 0.15,
                             test_ratio: float = 0.15):
    rnd = torch.rand(len(data.x))
    train_mask = [False if (x > val_ratio + test_ratio) else True for x in rnd]
    val_mask = [False if (val_ratio + test_ratio >= x) and (x > test_ratio) else True for x in rnd]
    test_mask = [False if (test_ratio >= x) else True for x in rnd]
    return torch.tensor(train_mask), torch.tensor(val_mask), torch.tensor(test_mask)
```

# GAT (Graph Attention Networks)

GATConv https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.conv.GATConv を使った GAT (Graph Attention Networks) クラスを実装します。


```python
import torch.nn.functional as F
from torch_geometric.nn import GATConv

class GAT(torch.nn.Module):
    def __init__(self):
        super(GAT, self).__init__()
        self.hid = 8
        self.in_head = 8
        self.out_head = 1
        
        
        self.conv1 = GATConv(dataset.num_features, self.hid, heads=self.in_head, dropout=0.6)
        self.conv2 = GATConv(self.hid*self.in_head, dataset.num_classes, concat=False,
                             heads=self.out_head, dropout=0.6)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index
                
        x = F.dropout(x, p=0.6, training=self.training)
        x = self.conv1(x, edge_index)
        x = F.elu(x)
        x = F.dropout(x, p=0.6, training=self.training)
        x = self.conv2(x, edge_index)
        
        return F.log_softmax(x, dim=1)
```

cuda が使えれば cuda を、使えなければ cpu を使うようにします。


```python
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device = "cpu"
```

## KarateClub データセットのノードラベル予測

データセットとモデルを次のようにします。


```python
dataset = datasets.KarateClub()
data = dataset[0].to(device)
model = GAT().to(device)
```

train, val, test データセットを次のようにセットします。



```python
train_mask, val_mask, test_mask = train_val_test_split(data)

data.train_mask = train_mask
data.val_mask = val_mask
data.test_mask = test_mask
```

次のようにして学習します。


```python
optimizer = torch.optim.Adam(model.parameters(), lr=0.005, weight_decay=5e-4)

losses = []
for epoch in range(2000):
    model.train()
    optimizer.zero_grad()
    out = model(data)
    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
    losses.append(loss.detach().numpy())

    if epoch % 200 == 0:
        print(loss)
    
    loss.backward()
    optimizer.step()
```

    tensor(1.4991, grad_fn=<NllLossBackward0>)
    tensor(0.7117, grad_fn=<NllLossBackward0>)
    tensor(0.9274, grad_fn=<NllLossBackward0>)
    tensor(0.4589, grad_fn=<NllLossBackward0>)
    tensor(0.4229, grad_fn=<NllLossBackward0>)
    tensor(0.4132, grad_fn=<NllLossBackward0>)
    tensor(0.1550, grad_fn=<NllLossBackward0>)
    tensor(0.7361, grad_fn=<NllLossBackward0>)
    tensor(0.5670, grad_fn=<NllLossBackward0>)
    tensor(0.3323, grad_fn=<NllLossBackward0>)


学習曲線は次のようになりました。


```python
import matplotlib.pyplot

plt.plot(losses, alpha=0.8)
```


![Qiita_PyG_(PyTorch_Geometric)_で_GAT_(Graph_Attention_Networks)_38_1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/211162/38debc8a-063b-f351-e2f1-b75c263d6cdd.png)







    

    


今回、目的変数 y は４種類のラベルになりますが、その予測結果は次のようになります。


```python
model(data)
```




    tensor([[-8.9359e+00, -1.1426e+00, -5.7290e+00, -3.8914e-01],
            [-6.4364e+00, -3.2092e-01, -1.4640e+00, -3.1798e+00],
            [-5.8728e+00, -3.6128e-01, -1.4878e+00, -2.5966e+00],
            [-7.9264e+00, -1.9967e-01, -1.9631e+00, -3.2134e+00],
            [-1.4734e+01, -5.6874e+00, -1.6281e+01, -3.3947e-03],
            [-7.4613e+00, -2.6100e+00, -5.5670e+00, -8.1134e-02],
            [-8.0405e+00, -2.9955e+00, -9.3368e+00, -5.1738e-02],
            [-9.7762e+00, -2.5239e-01, -1.5019e+00, -8.0933e+00],
            [-1.0885e+01, -8.3299e-02, -3.1077e+00, -3.3467e+00],
            [-1.4758e+00, -1.2810e+00, -1.5274e+00, -1.2853e+00],
            [-8.8328e+00, -2.0148e+00, -8.0043e+00, -1.4367e-01],
            [-1.4758e+00, -1.2810e+00, -1.5274e+00, -1.2853e+00],
            [-6.7221e+00, -1.1259e-01, -2.5291e+00, -3.6671e+00],
            [-6.8040e+00, -1.8940e-01, -1.9213e+00, -3.6878e+00],
            [-1.4758e+00, -1.2810e+00, -1.5274e+00, -1.2853e+00],
            [-1.4758e+00, -1.2810e+00, -1.5274e+00, -1.2853e+00],
            [-1.7997e+01, -7.4172e+00, -1.4998e+01, -6.0135e-04],
            [-1.4183e+00, -1.3070e+00, -1.4834e+00, -1.3456e+00],
            [-2.5639e-04, -1.7825e+01, -8.2728e+00, -1.3858e+01],
            [-2.5783e+00, -1.0899e+00, -6.1406e-01, -3.0638e+00],
            [-1.6390e-04, -1.8391e+01, -8.7232e+00, -1.3795e+01],
            [-2.4957e+00, -1.0830e+00, -6.0055e-01, -3.4909e+00],
            [-1.4758e+00, -1.2810e+00, -1.5274e+00, -1.2853e+00],
            [-3.1565e-01, -2.9938e+00, -1.5305e+00, -5.4833e+00],
            [-1.4758e+00, -1.2810e+00, -1.5274e+00, -1.2853e+00],
            [-1.4758e+00, -1.2810e+00, -1.5274e+00, -1.2853e+00],
            [-2.1341e-01, -6.6087e+00, -1.6809e+00, -5.3788e+00],
            [-4.2086e+00, -2.8116e+00, -7.8278e-02, -8.0508e+00],
            [-2.2204e+00, -1.5440e+00, -4.3808e-01, -3.4222e+00],
            [-6.4373e-06, -2.5067e+01, -1.1945e+01, -1.6719e+01],
            [-2.7569e+00, -1.2134e+00, -5.0864e-01, -3.2697e+00],
            [-1.1346e+01, -5.7865e-01, -8.5537e-01, -4.2540e+00],
            [-3.9678e-01, -4.7832e+00, -1.2040e+00, -3.9543e+00],
            [-1.4352e+00, -3.7600e+00, -3.2454e-01, -4.1481e+00]],
           grad_fn=<LogSoftmaxBackward0>)



学習済みのモデルの性能は次のように求まります。


```python
model.eval()
_, pred = model(data).max(dim=1)
correct = float(pred[data.test_mask].eq(data.y[data.test_mask]).sum().item())
acc = correct / data.test_mask.sum().item()
print('Accuracy: {:.4f}'.format(acc))
```

    Accuracy: 0.8333


## PubMed データセットのノードラベル予測

データセットとモデルを次のようにします。




```python
dataset = datasets.Planetoid(root=".", name="PubMed")
data = dataset[0].to(device)
model = GAT().to(device)
```

train, val, test データセットを次のようにセットします。


```python
train_mask, val_mask, test_mask = train_val_test_split(data)

data.train_mask=train_mask
data.val_mask=val_mask
data.test_mask=test_mask
```

次のようにして学習します。


```python
optimizer = torch.optim.Adam(model.parameters(), lr=0.005, weight_decay=5e-4)

losses = []
for epoch in range(1000):
    model.train()
    optimizer.zero_grad()
    out = model(data)
    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
    losses.append(loss.detach().numpy())

    if epoch % 200 == 0:
        print(loss)
    
    loss.backward()
    optimizer.step()
```

    tensor(1.0939, grad_fn=<NllLossBackward0>)
    tensor(0.6320, grad_fn=<NllLossBackward0>)
    tensor(0.6077, grad_fn=<NllLossBackward0>)
    tensor(0.6039, grad_fn=<NllLossBackward0>)
    tensor(0.6311, grad_fn=<NllLossBackward0>)
    tensor(0.6243, grad_fn=<NllLossBackward0>)
    tensor(0.6043, grad_fn=<NllLossBackward0>)
    tensor(0.6226, grad_fn=<NllLossBackward0>)
    tensor(0.6192, grad_fn=<NllLossBackward0>)
    tensor(0.6110, grad_fn=<NllLossBackward0>)


学習曲線は次のようになりました。


```python
import matplotlib.pyplot

plt.plot(losses, alpha=0.8)
```


![Qiita_PyG_(PyTorch_Geometric)_で_GAT_(Graph_Attention_Networks)_50_1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/211162/6ab7e945-3428-7e1d-96bf-8b09c1f095e6.png)


    


学習済みのモデルの性能は次のように求まります。


```python
model.eval()
_, pred = model(data).max(dim=1)
correct = float(pred[data.test_mask].eq(data.y[data.test_mask]).sum().item())
acc = correct / data.test_mask.sum().item()
print('Accuracy: {:.4f}'.format(acc))
```

    Accuracy: 0.8562


# おわり

また、そのうち PyG (PyTorch Geometric)  についていろいろ試してみたいと思います。