グラフカーネルやグラフニューラルネットワーク(GNN)の論文を読んでいると、TUDatasetというワードが登場します。初めてこの言葉に遭遇した人は何のことやら、と思うかもしれませんが、TUDatasetはその名前にある通り、データセットの名前となっています。
今回はTUDatasetについて、データセットの概要と簡単にこれらのデータセットにアクセスする方法について解説します。
TUDatasetとは
TUDatasetは、TU Dortmund Universityが2020年にICMLのworkshopで発表した論文「TUDataset: A collection of benchmark datasets for learning with graphs」で登場したデータセットで、主にグラフニューラルネットワークやグラフカーネルの研究向けに整備された、多数のグラフデータから構成されるデータセットです。
グラフデータの種類は数十以上あり、様々なジャンルのデータセットが取り揃えられています。
例えば以下のカテゴリのグラフデータが整備されています。
- 分子構造
- ソーシャルネットワーク
- バイオインフォマティクス
- コンピュータビジョン
非常多くのデータセットがあることから、今後広く利用されることになると思います。
TUDatasetを利用する方法
データセットは色々なツールから利用することができ、例えば公式ホームページのDatasetページから、zipを直接落としてくることで利用することができますが、pytorchのグラフ用ライブラリであるpytorch geomatric(PyG)を利用すると、非常に簡単にこれらのデータセットを扱うことができます。
今回は、PyGを用いてTUDatasetにアクセスしてみましょう。
Pytorch geometricでTUDatasetを扱う
pytorch geometricがインストールしていない人は、こちらの記事に従ってインストールをしてください。
今回は、188個の分子構造のデータからなるMUTAGデータセットを入手してみましょう。
まずは、必要なライブラリをインポートします。
import matplotlib.pyplot as plt
import networkx as nx
import torchvision
import torchvision.transforms as transforms
from torch_geometric.utils import to_networkx
from torch_geometric.data import Data
続いて、データセットをダウンロードしてきます。188個のデータがあることがわかります。
from torch_geometric.datasets import TUDataset
dataset = TUDataset(root="./MUTAG", name="MUTAG")
len(dataset) # => 188
そのうちの、MUTAGデータセットに入っている1つの分子構造を見てみましょう。
data = dataset[0]
print(data) #=> Data(edge_index=[2, 38], x=[17, 7], edge_attr=[38, 4], y=[1])
17個のノードと38個のエッジとから構成され、各ノードが7次元の特徴量を持っていることがわかります。この辺りのattributeの見方は、pytorch geometricの解説記事をご覧ください。
最後に可視化してみましょう。可視化には、networkXを利用します。
import networkx as nx
from torch_geometric.utils import to_networkx
nxg = to_networkx(data)
nx.draw(nxg, with_labels=True)

このような分子構造をしていることがわかりました。
- TUDataset
- TUDataset: A collection of benchmark datasets for learning with graphs