【KGE/知識グラフ埋め込み】PyKEENを用いたKGE (Knowledge Graph Embedding)モデル構築ガイド

はじめに

本記事はPyKEENを使ってKGEモデルを構築したいが，ライブラリの全体像がつかめていないため何から手をつければ良いかわからない人向けの内容です．

公式ドキュメントが親切に書かれているので，本記事ではインストールからモデル構築&評価までに最低限必要なことを公式ドキュメントを参照しながら解説します．

筆者が大学の研究でPyKEENを用いてKGEモデルを構築した経験から得た知見をまとめています．私自身もKGEの研究歴は浅いですが，本記事の執筆時点でQiitaには"PyKEEN"にヒットする記事が無く，初心者の足がかりになるような記事があれば良いなと思ったので投稿します．

インストール

公式ドキュメントのInstallation通りに実行すれば問題ないです．以下にMacOS, Linuxユーザ向けのコマンドを示します．安定版と最新版でコマンドが異なります．

stable

pip install pykeen

latest

pip install git+https://github.com/pykeen/pykeen.git

MacOS, Linux以外の環境を使用する場合は以下のドキュメントを参照しpip installを実行しましょう．

トリプルについて

モデル構築に使用するトリプルは大きく分けて以下のどちらかになります．

1. PyKEENであらかじめ用意されたトリプル
2. 自分で用意したトリプル

PyKEENであらかじめ用意されたトリプルを用いる場合

モデル構築スクリプトの実行前に準備する必要はありません．次章のモデル構築時に参照するドキュメントの例に倣えば良いです．どのようなデータセットが準備されているか知りたい方は以下のDatasetsのページを参照してください．

自分で用意したトリプルを用いる場合

手順を以下に示します．

1. tsvファイル or numpy.ndarray型でトリプルを準備
2. TriplesFactoryオブジェクトを生成（from_pathやfrom_labeled_triples等のクラスメソッドを使用）
3. splitメソッドを使用して訓練データ，検証データ，テストデータに分割

カラムの順序
tsvファイルでもndarrayでも，カラムがhead, relation, tailの順番になるように作成してください．

tsvの場合

my_triples.tsv

ONE PIECE    is_written_by     Eiichiro Oda
Attack on Titan    is_written_by    Hajime Isayama
...

from pykeen.triples import TriplesFactory
from pykeen.pipeline import pipeline

tf = TriplesFactory.from_path("my_triples.tsv")
train, valid, test = tf.split([0.8, 0.1, 0.1]) # 8:1:1に分割
result = pipeline(
    training=train,
    validation=valid,
    testing=testing,
    model='TransE',
    stopper='early',
    epochs=5,
)
result.save_to_directory('test_result')

ndarrayの場合

from pykeen.triples import TriplesFactory
from pykeen.pipeline import pipeline
import numpy as np

# triplesはndarrayでトリプルデータが格納されていると仮定します．
tf = TriplesFactory.from_labeled_triples(triples)
train, valid, test = tf.split([0.8, 0.1, 0.1]) # 8:1:1に分割
result = pipeline(
    training=train,
    validation=valid,
    testing=testing,
    model='TransE',
    stopper='early',
    epochs=5,
)
result.save_to_directory('test_result')

モデル構築

公式ドキュメントのFirst Stepsを参照します．

"Training a Model"と"Deeper Configuration"節にpykeen.pipeline.pipeline()を利用した以下のような例がたくさん書かれています．

from pykeen.pipeline import pipeline
pipeline_result = pipeline(
    dataset='Nations',
    model='TransE',
)
pipeline_result.save_to_directory('nations_transe')

それぞれの例の違いはpipeline()に与える引数や与え方です．

これらの具体例を見て大雑把に使い方を理解し，pykeen.pipeline.pipeline()のページを見て学習に設定可能なパラメータを体系的に理解するのが一般的な流れです．

補足：具体例がわからなくても大丈夫

具体例を見た時点で「学習のために値渡してるのはわかるけど，そもそも学習にどう影響する引数なのかが分からん」とか「この値どういうこと？」という疑問が発生する方がいるかと思います（筆者がそうでした）．その原因として以下の二つが考えられます．

1. KGEに対しての理解がまだ深まっていない
2. オブジェクトの仕様を知らない

前者は自分が実装したいKGEモデルの論文を読み込む，後者はpipeline()関数のページで各引数の型や説明を見たり詳細ページに飛んで解説を見るなどで対処できるかと思います．また逆に，pipeline()関数のページで得たパラメータの知識を活用してKGEの論文をより理解するというのもアリだと思います．

パイパーパラメータチューニング

pykeen.hpo.hpo_pipeline()を使うと簡単にハイパーパラメータチューニングができます．デフォルトの最適化手法はベイズ最適化で，探索範囲や探索候補を設定することが可能です．

プログラムの書き方は各パラメータの型ごとに異なりますが，pipelineのxxx_kwargs引数にdictの要素として渡すパラメータをhpo_pipelineのxxx_kwargs_ranges引数に同様に渡すという点は共通しています．

以下に具体例を示します．

categorical

from pykeen.hpo import hpo_pipeline
hpo_pipeline_result = hpo_pipeline(
    dataset='Nations',
    model='KG2E',
    model_kwargs_ranges=dict(
        dist_similarity=dict(type='categorical', choices=['KL', 'EL']),
    ),
)

カテゴリカルなパラメータはchoices引数に探索候補をリスト型で渡します．

bool

from pykeen.hpo import hpo_pipeline
hpo_pipeline_result = hpo_pipeline(
    dataset='Nations',
    model='TransE',
    training_loop='sLCWA',
    negative_sampler_kwargs_ranges=dict(
        filtered=dict(type=boolean),
    ),
)

真偽値型のパラメータはtype以外に指定する必要はありません．内部的にchoices=[True, False]としたカテゴリカル変数と捉えられます．

int, float

上限と下限を指定した上で線形，指数，対数などのスケールやステップ幅などを調整できます．詳しくはドキュメントを参照してください．

評価

pipelineを実行するとPipelineResultクラスのオブジェクトが返されます．First stepsにもあるように，学習に関する全ての情報を出力する場合はsave_to_directoryを実行するのがおすすめです．

保存先ディレクトリの構造を以下に示します．

保存先ディレクトリの構造

.
├── metadata.json
├── results.json
├── trained_model.pkl
└── training_triples
    ├── base.pth
    ├── entity_to_id.tsv.gz
    ├── numeric_triples.tsv.gz
    └── relation_to_id.tsv.gz

ファイル名	説明
metadata.json	モデル情報やパラメータ
results.json	loss関数の値や評価指標の値
trained_model.pkl	構築したモデルのpickleファイル
base.pth	トリプルに関する情報(エンティティ・リレーション総数，inverse_triple)
entity_to_id.tsv.gz	エンティティラベルと内部id(int型)の対応
relation_to_id.tsv.gz	リレーションラベルと内部id(int型)の対応
numeric_triples.tsv.gz	トリプルを内部idに置換したもの

論文等でよく目にするHits@k, MRR(InverseHarmonicMeanRank)などはresults.jsonで確認しましょう．評価指標クラスは以下のページにまとめられています．

おわりに

本記事では初心者に向けたPyKEENでのKGEモデル構築方法を解説しました．

質問・意見等あれば気軽にコメントしていただけると幸いです！

参考