Zenn
Open2

DGLとなかよく

Python
#
gnn
TsumuraTsumura

Node Classification

https://docs.dgl.ai/tutorials/blitz/1_introduction.html#sphx-glr-tutorials-blitz-1-introduction-py
Node Classification(ノード分類): グラフ構造におけるノード(頂点)のカテゴリや属性を予測する問題

Loading dataset

今回はCoraデータセットという、論文をノード、引用関係をエッジとするネットワークを用いる。タスクは特定の論文のカテゴリを予測する。

import os

# 用いる深層学習フレームワークを指定
os.environ["DGLBACKEND"] = "pytorch"

import dgl
import dgl.data
import torch
import torch.nn as nn
import torch.nn.functional as F

dataset = dgl.data.CoraGraphDataset()

DGLのグラフでは、ノード特徴量とエッジ特徴量を、ndataとedataという辞書のような属性に格納できる。DGLのCoraデータセットのグラフには、以下のようなノード特徴量が含まれている。

  • train_mask: ノードが訓練セットに含まれるかどうかを示すブーリアンテンソル
  • val_mask: ノードが検証セットに含まれるかどうかを示すブーリアンテンソル
  • test_mask: ノードがテストセットに含まれるかどうかを示すブーリアンテンソル
  • label: ノードの正解カテゴリ
  • feat: ノード特徴量

Defining a Graph Convolutional Network (GCN)

今回は二層のGCNを設計する。設計には torch.nn.Module を継承すればOK(簡単!)

from dgl.nn import GraphConv


class GCN(nn.Module):
    def __init__(self, in_feats, h_feats, num_classes):
        super(GCN, self).__init__()
        self.conv1 = GraphConv(in_feats, h_feats)
        self.conv2 = GraphConv(h_feats, num_classes)

    def forward(self, g, in_feat):
        h = self.conv1(g, in_feat)
        h = F.relu(h)
        h = self.conv2(g, h)
        return h


# Create the model with given dimensions
model = GCN(g.ndata["feat"].shape[1], 16, dataset.num_classes)

Training the GCN

pytorch の書き方がわかっていれば楽勝

def train(g, model):
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    best_val_acc = 0
    best_test_acc = 0

    # ノード分類なので `ndata`から取得
    features = g.ndata["feat"]
    labels = g.ndata["label"]
    train_mask = g.ndata["train_mask"]
    val_mask = g.ndata["val_mask"]
    test_mask = g.ndata["test_mask"]
    for e in range(100):
        # Forward
        logits = model(g, features)

        # Compute prediction
        pred = logits.argmax(1)

        # Compute loss
        # Note that you should only compute the losses of the nodes in the training set.
        loss = F.cross_entropy(logits[train_mask], labels[train_mask])

        # Compute accuracy on training/validation/test
        train_acc = (pred[train_mask] == labels[train_mask]).float().mean()
        val_acc = (pred[val_mask] == labels[val_mask]).float().mean()
        test_acc = (pred[test_mask] == labels[test_mask]).float().mean()

        # Save the best validation accuracy and the corresponding test accuracy.
        if best_val_acc < val_acc:
            best_val_acc = val_acc
            best_test_acc = test_acc

        # Backward
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if e % 5 == 0:
            print(
                f"In epoch {e}, loss: {loss:.3f}, val acc: {val_acc:.3f} (best {best_val_acc:.3f}), test acc: {test_acc:.3f} (best {best_test_acc:.3f})"
            )


model = GCN(g.ndata["feat"].shape[1], 16, dataset.num_classes)
train(g, model)