こんにちは。シンプルフォーム株式会社 にてインフラエンジニアをしています、山岸です。
当社では Amazon Neptune というグラフデータベースを用いた機能開発を行っています。以前、RDS データベース由来のデータを Neptune に流し込むための ETL パイプライン実装に関する記事を執筆しました。
本記事では、以前あまり触れられなかった Neptune ML について実際に触ってみたため、その概要や使い方についてご紹介したいと思います。
概要
- シンプルフォームでは、グラフ DB を用いた機能を提供しています。グラフ DB の基盤として Amazon Neptune を使用しています。グラフの特性を活かしたより高度な機能提供に向けて、Graph Neural Network (GNN) のような技術についても価値検証を行なっています。
- Neptune には、グラフデータに対する機械学習モデルを構築・運用するための Neptune ML という機能が備わっています。本記事では、レコメンデーションシステムの評価などによく用いられる MovieLens というデータセットを使用して、Neptune ML による GNN モデルの構築 〜 デプロイの一連のプロセスについて解説します。
本記事で触れないこと
- Neptune ML の使い方の解説に主眼を置いているため、グラフ DB そのものや、GNN の数学的な理論については扱いません。
- Neptune ML では機械学習モデルとして「Graph Neural Network (GNN) モデル」と「Knowledge Graph Embedding (KGE) モデル」がサポートされています。本記事では扱うのは GNN モデルであり、KGE モデルについては扱いません。
- 実際的なユースケースとしては、時間の経過とともにグラフの状態が変化することが多いかと思いますが、今回は簡単のため静的なグラフを想定します。増分データに対するモデル変換や推論エンドポイントの更新については、別の機会に譲りたいと思います。
前提知識
本題に入る前に、GNN と Neptune ML について少しだけ触れておきたいと思います。
GNN について
Graph Neural Network (GNN) とは、一言でいえば「深層学習の技術をグラフ領域に拡張したもの」です。
例えば画像などは深層学習が扱う対象として代表的なものですが、画像は各画素をノードとし、隣接する画素同士がエッジで結ばれたグラフであるとも解釈することが可能です。画像には二次元空間上の制約が存在するわけですが、そのような制約を取り払ってより一般化したデータ構造がグラフであると捉えることができます。実際、画像領域でいうところの畳み込みニューラルネットワーク (CNN) に相当するものがグラフにおいても考案されており、これは Graph Convolutional Network (GCN) と呼ばれます。
より一般化されたデータ構造であるという意味で、GNN のユースケースは多岐にわたります。
- レコメンデーション
- ソーシャルネットワーク分析
- 化合物の物性推定 ／ 製薬
- 交通や物流の予測・最適化 ... etc.
Neptune ML について
Neptune ML は、Amazon Neptune がもつ GNN モデル学習・推論のための機能です。
機械学習マネージドサービスである Amazon SageMaker と統合されており、また GNN モデル構築のための Python ライブラリである Deep Graph Library が内部的に使用されています。これにより、GNN モデル開発に関する専門知識をもたないエンジニアであっても、短期間でモデルの構築〜デプロイ、推論までを行うことができます。
Neptune ML では、以下の推論タスクがサポートされています。[1]
デモ概要
デモ用のデータセットとして MovieLens を使用します。このデータセットには、ユーザー情報、映画情報、および映画に対するユーザーの評価情報が含まれており、レコメンデーションシステムの評価などによく利用されます。
グラフ上では、ユーザーを User ノード、映画を Movie ノード、評価を Rate エッジとして表現できそうです。Rate エッジのプロパティとして評価スコアを持たせ、評価スコアやその他の特徴量を学習したグラフネットワークを構築することで、実在しない Rate エッジの評価スコアを推論値として算出できます。
あるユーザーがまだ視聴していない映画について評価スコアを推論し、そのスコアが高ければそのユーザーに対してレコメンドするといった活用が考えられますので、これを以てレコメンデーションということにしたいと思います。
参考
MovieLens を用いたエッジ回帰推論については、以下にサンプルのノートブックがあります。後述の特徴量設計もこちらを参考にしたいと思います。
こちらのノートブックの Note にも記載されている通り、予測精度の良いモデルを生成することを目的としていないことに留意します。
Note: The configuration used in this notebook specifies only a minimal set of configuration options meaning that our model's predictions are not as accurate as they could be. The parameters included in this configuration are one of a couple of sets of options available to the end user to tune the model and optimize the accuracy of the resulting predictions.
実施手順
基本的には 公式ドキュメント に記載の手順に沿って進めていきます。
事前準備
リソースの作成
今回の検証に使用する Neptune クラスターを作成します。検証目的のため、db.t4g.medium の小さめのインスタンスサイズ、かつ 1 台構成（リードレプリカなし）で構築しました。
また、Neptune 用と SageMaker 用でそれぞれ IAM ロールを作成します。これらは構築した Neptune クラスターに関連づけられている必要があります。
データ投入
MovieLens データセット (ml-1m) に含まれる movies.dat, users.dat, および ratings.dat を使用し、以下のようなスキーマを持つグラフを構築します。
今回は S3 に配置した元データの DAT ファイルを Glue ジョブで Gremlin 形式 に整形し、このデータを Neptune Bulk Loader を使用して Neptune クラスターに対して投入しました。Neptune Bulk Loader については 以前の記事 で詳しく紹介しているので、よろしければ併せてご覧ください。
参考までに、投入したデータの一部を以下に示します。
投入した Gremlin 形式データの一部
頂点データ
~id,~label,gender:String(single),age:Int(single),occupation:String(single),zip_code:String(single)
user_56,user,M,35,20,60440
user_73,user,M,18,4,53706
user_110,user,M,25,2,90803
...
~id,~label,title:String(single),genres:String(single)
movie_1,movie,Toy Story (1995),"Animation,Children's,Comedy"
movie_2,movie,Jumanji (1995),"Adventure,Children's,Fantasy"
movie_3,movie,Grumpier Old Men (1995),"Comedy,Romance"
...
エッジデータ
~id,~label,~from,~to,score:Double(single),timestamp:String(single)
4cf9d390697843fa3003357b78839c36,rate,user_1,movie_1193,5.0,978300760
cefe1a63207eb47c211ebb4c26076a2b,rate,user_1,movie_661,3.0,978302109
dc5d6ec8f6351098fc6efdfb1765bfc0,rate,user_1,movie_914,3.0,978301968
...
クォータ上限の引き上げ
アカウントのデフォルトの状態だと、SageMaker Processing ジョブを生成するステップでクォータ上限を超過する可能性があります。必要に応じて適宜クォータ上限引き上げをリクエストしてください。筆者の環境では以下のリクエストを作成し、本記事の内容を最後まで実施できました。
- [Amazon SageMaker] ml.r5.large for processing job usage ...
2 - [Amazon SageMaker] ml.g4dn.2xlarge for training job usage ...
4
Step-1. データエクスポート
Neptune クラスターに存在するデータを、GNN モデル学習用に S3 にエクスポートします。データエクスポートのための手段はいくつか存在します [2] が、本記事では Neptune-Export Service を使用する方法について扱います。
1-1. Neptune-Export Service のデプロイ【初回のみ】
こちらの手順 にあるリージョン一覧の中から使用するリージョンに該当する CloudFormation スタックを選択してデプロイします。スタック作成の際、パラメータの [Network Configuration] にて Neptune クラスターが稼働している VPC、およびそのプライベートサブネットを指定します。
ネストを含む全てのスタック作成が完了すると、以下のようなアーキテクチャで各種リソースがデプロイされます。具体的には、API Gateway Private API、Lambda 関数、Batch ジョブ（正確にはジョブ定義とジョブキュー）、およびこれらに付随するリソースなどがデプロイされます。
1-2. ターゲット定義
推論タスクの問題設定に合わせてターゲットを定義します。今回は rate エッジのプロパティである score に対するエッジ回帰推論であるため、以下のようになります。"split_rate" には、学習用・検証用・評価用として使用するデータの割合を表す３つの小数のリストを渡します。
targets = [
{
"edge":["user", "rate", "movie"],
"property":"score",
"type":"regression",
"split_rate":[0.8, 0.1, 0.1]
}
]
1-3. 特徴量設計
モデル学習に使用する特徴量について、データの処理方法とともに設計します。今回は、サンプルノートブックでも使用されている title (word2vec), age (bucket_numerical) に gender (category) を加えた３つの特徴量を使用してみたいと思います。
features = [
{
"node":"movie",
"property":"title",
"type":"text_word2vec"
},
{
"node":"user",
"property":"age",
"type":"bucket_numerical",
"range":[1, 100],
"num_buckets":10,
"imputer":"mean"
},
{
"node":"user",
"property":"gender",
"type":"category"
}
]
特徴量のデータ処理方法に関する詳細は、以下の公式ドキュメントをご確認ください。
1-4. エクスポート実行
Neptune-Export Service の API Gateway エンドポイントに対して、以下のような POST リクエストを送ります。これにより、バックエンドの Lambda 関数経由で Batch ジョブが起動されます。Batch ジョブの実行が完了すると、"outputS3Path" で指定した S3 パス配下にディレクトリが作成され、出力が配置されます。
NEPTUNE_EXPORT_API_URI = "https://xxxxxxxx.execute-api.ap-northeast-1.amazonaws.com/v1/neptune-export"
data = {
"command": "export-pg",
"outputS3Path": f"s3://{BUCKET_NAME}/movielens/exported-data",
"params": {
"endpoint": NEPTUNE_CLUSTER_ENDPOINT,
"profile": "neptune_ml",
},
"additionalParams":{
"neptune_ml":{
"version":"v2.0",
"targets":targets,
"features":features,
}
}
}
response = requests.post(
url=NEPTUNE_EXPORT_API_URI,
headers={"Content-Type": "application/json"},
json=data,
)
content = json.loads(response.content.decode("utf-8"))
Step-2. データ処理
このステップでは、エクスポートされたデータをもとに、モデル学習中に Deep Graph Library (DGL) が使用する情報を生成します。
2-1. training-data-configuration.json の編集【任意】
Step-1 のデータエクスポート先の S3 パスを確認してみると、以下のようなディレクトリ構成になっています。edges/ , nodes/ 配下には実データを含んでおり、その他の JSON ファイルはメタデータとなっています。
exported-data/
└── yyyymmdd_hhMMss/
├── config.json
├── edges/
├── nodes/
├── stats.json
└── training-data-configuration.json
このうち、training-data-configuration.json は後続ステップにおける処理方法を記述しています。例えば、学習に使用する特徴量の前処理方法などです。これらの情報はデータエクスポート時のリクエスト内容に由来します。
training-data-configuration.json のサンプル
{
"version":"v2.0",
"query_engine":"gremlin",
"graph":{
"nodes":[
{
"file_name":"nodes/movie.consolidated.csv",
"separator":",",
"node":["~id", "movie"],
"features":[
{
"feature":["title", "title", "text_word2vec"],
"language":["en_core_web_lg"]
},
{
"feature":["genres", "genres", "auto"]
}
]
},
{
"file_name":"nodes/user.consolidated.csv",
"separator":",",
"node":["~id", "user"],
"features":[
{
"feature":["age", "age", "bucket_numerical"],
"range":[1, 100],
"bucket_cnt":10,
"slide_window_size":0,
"imputer":"mean"
},
{
"feature":["occupation", "occupation", "auto"]
},
{
"feature":["zip_code", "zip_code", "auto"]
},
{
"feature":["gender", "gender", "category"]
},
{
"feature":["avg_score", "avg_score", "numerical"],
"norm":"min-max",
"imputer":"mean"
}
]
}
],
"edges":[
{
"file_name":"edges/%28user%29-rate-%28movie%29.consolidated.csv",
"separator":",",
"source":["~from", "user"],
"relation":["", "rate"],
"dest":["~to", "movie"],
"features":[
{
"feature":["timestamp", "timestamp", "auto"]
}
],
"labels":[
{
"label":["score", "regression"],
"split_rate":[0.8, 0.1, 0.1]
}
]
}
]
},
"warnings":[]
}
グラフ上のデータに変更がない限り、実データ部分のエクスポート結果は変わりません。training-data-configuration.json の内容を編集して既存のものと置き換えることで、データエクスポートの処理を再度実行せずとも、後続ステップでの処理方法を変更することができます。
2-2. DataProcessing ジョブの実行
データ処理用に提供されている /ml/dataprocessing エンドポイントに POST リクエストを送ることで、DataProcessing ジョブを開始できます。レスポンスボディには DataProcessing ジョブ ID が含まれ、これは後続のモデルトレーニングで使用します。
PROCESSING_ENDPOINT_URL = f"https://{NEPTUNE_CLUSTER_ENDPOINT}:8182/ml/dataprocessing"
data = {
"inputDataS3Location": f"s3://{BUCKET_NAME}/movielens/exported-data/{timestamp}/",
"processedDataS3Location": f"s3://{BUCKET_NAME}/movielens/processed-data/",
"sagemakerIamRoleArn": NEPTUNE_SAGEMAKER_ROLE_ARN,
"neptuneIamRoleArn": NEPTUNE_CLUSTER_ROLE_ARN,
}
response = requests.post(
url=PROCESSING_ENDPOINT_URL,
headers={"Content-Type": "application/json"},
json=data,
)
content = json.loads(response.content.decode("utf-8"))
print(content)
# {"id": "01c4c6b2-b898-4d28-841e-b98e55395d15"}
Step-3. モデルトレーニング
このステップでは GNN モデルの学習を行います。
3-1. model-hpo-configuration.json の編集【任意】
Step-2 データ処理出力先の S3 パスを確認してみると、以下のようなディレクトリ構成になっています。ノードやエッジなどの実データは graph.bin に含まれています。
01c4c6b2-preloading.../
└── preloading-output/
├── features.json
├── graph.bin
├── graph.info
├── info.pkl
├── model-hpo-configuration.json
├── train_instance_recommendation.json
└── updated_export_config.json
出力されたファイル一覧の中に model-hpo-configuration.json というファイルが存在します。このファイルには、モデル調整時に使用されるハイパーパラメータが含まれます。ファイルを編集して既存のものと置き換えることで、モデル学習時のハイパーパラメータを変更することができます。
model-hpo-configuration.json のサンプル
{
"models":[
{
"model":"rgcn",
"task_type":"edge_regression",
"eval_metric":{"metric":"mse"},
"eval_frequency":{
"type":"evaluate_every_epoch",
"value":1
},
"1-tier-param":[...],
"2-tier-param":[...],
"3-tier-param":[...],
"fixed-param":[...]
}
]
}
3-2. ModelTrainingJob の実行
モデルトレーニング用に提供されている /ml/modeltraining エンドポイントに POST リクエストを送ることで、ModelTraining ジョブを開始できます。レスポンスボディには ModelTraining ジョブ ID が含まれ、これは後続の推論エンドポイントデプロイで使用します。
TRAINING_ENDPOINT_URL = f"https://{NEPTUNE_CLUSTER_ENDPOINT}:8182/ml/modeltraining"
data = {
"dataProcessingJobId": "01c4c6b2-b898-4d28-841e-b98e55395d15",
"trainModelS3Location": f"s3://{BUCKET_NAME}/movielens/model-artifacts/",
"modelName": "rgcn",
"trainingInstanceType": "ml.g4dn.2xlarge",
"sagemakerIamRoleArn": NEPTUNE_SAGEMAKER_ROLE_ARN,
"neptuneIamRoleArn": NEPTUNE_CLUSTER_ROLE_ARN,
}
response = requests.post(
url=TRAINING_ENDPOINT_URL,
headers={"Content-Type": "application/json"},
json=data,
)
content = json.loads(response.content.decode("utf-8"))
print(content)
# {"id": "30714daf-2494-4f42-bc5f-c6b49a70a619"}
モデル学習が完了すると、trainModelS3Location で指定した S3 パスにモデルアーティファクトのファイル群が出力されます。
Step-4. 推論エンドポイントデプロイ
推論エンドポイント用に提供されている /ml/endpoints エンドポイントに POST リクエストを送ることで、SageMaker 推論エンドポイントが作成されます。
ML_ENDPOINTS_ENDPOINT_URL = f"https://{NEPTUNE_CLUSTER_ENDPOINT}:8182/ml/endpoints"
data = {
"mlModelTrainingJobId": "30714daf-2494-4f42-bc5f-c6b49a70a619",
"update": False,
"instanceType": "ml.m5.large",
"instanceCount": 1,
"neptuneIamRoleArn": NEPTUNE_CLUSTER_ROLE_ARN,
}
response = requests.post(
url=ML_ENDPOINTS_ENDPOINT_URL,
headers={"Content-Type": "application/json"},
json=data,
)
content = json.loads(response.content.decode("utf-8"))
# {"id": "d64fd143-ddf3-4126-8fb2-27505f1e56b4"}
推論エンドポイントは、SageMaker マネジメントコンソールの [推論] - [エンドポイント] から確認できます。エンドポイント名は、作成時のリクエスト ID の一部を含むものになります。ステータスが [InService] になったら、推論できる状態になっています。
Step-5. 推論
作成された推論エンドポイントに対して、実際にエッジ回帰推論を行なっていきます。
5-1. クエリ記述方法
クエリの起点となる GraphTraversal オブジェクト g を生成します。
from gremlin_python.structure.graph import Graph
from gremlin_python.process.graph_traversal import __
from gremlin_python.process.traversal import T
from gremlin_python.driver.driver_remote_connection import DriverRemoteConnection
from gremlin_python.driver.aiohttp.transport import AiohttpTransport
graph = Graph()
conn = DriverRemoteConnection(
f"wss://{NEPTUNE_CLUSTER_ENDPOINT}:8182/gremlin",
"g",
transport_factory=lambda: AiohttpTransport(ssl_options=ssl_opts),
)
g = graph.traversal().with_remote(conn)
Score プロパティの値を推論したいエッジを作成します。（generate_edge_id() はエッジ ID を生成するための関数として記載しています。適宜自身のロジックに読み替えてください）
rate_edge_id = generate_edge_id(user_id, movie_id)
query = g\
.V(user_id).add_e("rate").to(__.V(movie_id))\
.property(T.id, rate_edge_id)
query.iterate()
以下が推論本体のコードになります。Neptune#ml. から始まる Neptune 独自の With ステップをいくつか含みます。推論エンドポイント名や SageMaker 用に作成した IAM ロールを指定します。また、今回は回帰タスクなので、推論結果を算出する際の With ステップには Neptune#ml.regression を指定します。
inference_endpoint_name = "d64fd143-2024-02-29-08-57-7350000-endpoint"
y_pred = g\
.with_("Neptune#ml.endpoint", inference_endpoint_name)\
.with_("Neptune#ml.iamRoleArn", NEPTUNE_SAGEMAKER_ROLE_ARN)\
.E(rate_edge_id).properties("score")\
.with_("Neptune#ml.regression")\
.value().to_list()
predicted_score = y_pred[0]
推論用に作成したエッジは、グラフ上には実際には存在していない想定なので、一連の処理実行が完了したら削除しておきます。
g.E(rate_edge_id).drop().iterate()
5-2. 評価用データに対する推論結果
評価用データとして事前に実データからランダム抽出しておいたエッジデータについて、「実際の評価スコア（＝実際スコア）」と「予測された評価スコア（＝予測スコア）」を比較してみたいと思います。評価用エッジデータは 1,200 件あまりで、全体の 0.1% 強のデータ量です。（勿論、グラフには投入しておらず学習にも使用していません）
前述の推論処理を、評価用データフレーム全体に対して実行しました。以下は、実際スコアと予測スコアをそれぞれ軸に取ったバイオリンプロットです。
後片付け
検証用途で作成した、今後不要なリソースは削除しておきましょう。特に Neptune クラスターや SageMaker 推論エンドポイントなどは、稼働しているとそれなりのコストになってくるので注意が必要です。
さいごに
Neptune ML の使い方についてご紹介しました。いかがでしたでしょうか。
モデル開発の専門知識がなくても、グラフ ML の一連のプロセスを実施できるのは非常に便利だと感じました。一方、今回のようなモデル変換を伴わない比較的簡単なケースにおいても手順が多く、インフラ面での初見ハードルはやや高いようにも感じます。Neptune ML 導入検討の際の参考になれば幸いです。
また、検証にあたり AWS サポート様には多大なるご助力を頂きました。この場をお借りして感謝を申し上げたいと思います。
リアルタイム法人調査システム「SimpleCheck」を開発・運営するシンプルフォーム株式会社の開発チームのメンバーが、日々の開発で得た知見や試してみた技術などについて発信していきます。 Publication 運用への移行前の記事は zenn.dev/simpleform からご覧ください。
Discussion