MINDデータセットを利用して、ニュース推薦モデル(NAML, NRMS)をカスタマイズしつつ学習するためのコード群が置かれています。
linux上のdockerで動作させる想定です。以下の環境で動作確認済みです。
- linux (Ubuntu 20.04 LTS)
- docker: version 20.10.6, build 370c289
- docker-compose: 1.29.1, build c34c88b2
- nvidia-container-toolkit: 1.5.1-1 amd64
- GPU: NVIDIA GeForce RTX 2080 Ti
git clone --recursive https://github.com/stockmarkteam/mind-recommenders-pytorch
mv mind-recommenders-pytorch
cp .env.sample .env
.env
に定義された環境変数は以下のとおりです。必要に応じて変更可能ですが、以降の説明はデフォルト設定を前提として行われます。
COMPOSE_PROJECT_NAME
:- docker-composeの環境変数。詳細はこちら。
DEVICE
(デフォルト値:gpu
):- dockerで利用するデバイスを指定します。
gpu, cpu
のうちいずれかを選択してください。一応切り替えができるようになっていますが、cpu
設定での前処理/学習スクリプトの動作は未確認です。
- dockerで利用するデバイスを指定します。
DATASET_PATH
(デフォルト値:$(PWD)/dataset
):- mind datasetを保存するhostディレクトリ。container上では
dataset/
にmountされます。
- mind datasetを保存するhostディレクトリ。container上では
MODEL_PATH
(デフォルト値:$(PWD)/models
):- GloVe, Transformerのpretrained modelを保存するhostディレクトリ。container上では
models/
にmountされます。
- GloVe, Transformerのpretrained modelを保存するhostディレクトリ。container上では
LOG_PATH
(デフォルト値:$(PWD)/logs
)::- 学習のログを保存するhostディレクトリ。container上では
logs/
にmountされます。
- 学習のログを保存するhostディレクトリ。container上では
VENV_PATH
:- pythonの仮想環境をinstallするhostディレクトリ。container上では
.venv/
にmountされます。
- pythonの仮想環境をinstallするhostディレクトリ。container上では
JUPYTER_PORT
:- container上で立ち上げたjupyter notebookにhostOS上のブラウザからアクセスするためbindするportを指定します。(default:
8888
)
- container上で立ち上げたjupyter notebookにhostOS上のブラウザからアクセスするためbindするportを指定します。(default:
TENSORBOARD_PORT
:- container上で立ち上げたtensorboardにhostOS上のブラウザからアクセスするためbindするportを指定します。(default:
6006
)
- container上で立ち上げたtensorboardにhostOS上のブラウザからアクセスするためbindするportを指定します。(default:
make setup
公式サイトから訓練データセット・開発データセットのzipファイルをDLして、containerから見える場所に配置してください。 迷ったらこのREADMEと同じディレクトリに配置すれば問題ありません。
make sh
コンテナ内で以下のコマンドを実行することにより、必要な前処理が全て行われます。
pipenv run preprocess-all data_path.train_zip=<MINDxxx_train.zipのpath> data_path.valid_zip=<MINDxxx_dev.zipへのpath>
通常版データセットを利用する場合は、上記コマンドの引数にparams.dataset_type=large
を追加してください。
ここで行われる各処理の概要ついては、こちらをご確認ください。
コンテナ内で以下のコマンドを実行することにより、モデルが学習できます。
pipenv run train
当社のブログ記事で言及した12通りのモデルをすべて学習したい場合は、以下のコマンドを実行してください。
pipenv run train-all
指定できるオプションの一例は以下のとおりです。
model
:naml
ornrms
(default:nrms
)
embedding_layer
:word_embedding
ortransformer
(default: word_embedding)
hparams.article_attributes
:- 利用する記事属性を
[title,body,category,subcategory]
から指定(default:[title,body,category,subcategory]
)
- 利用する記事属性を
hparams.n_epochs
:- 訓練のエポック数
hparams.max_title_length
:- 最大タイトルトークン長(default:
30
)
- 最大タイトルトークン長(default:
hparams.max_body_length
:- 最大本文トークン長(default:
128
)
- 最大本文トークン長(default:
hparams.batch_size.train
:- train datasetのbatch size(default: 利用するembedding layerに応じて変化)
hparams.batch_size.valid
:- validation datasetのbatch size(default: 利用するembedding layerに応じて変化)
hparams.accumulate_grad_batches
:- この数値と同じstep数が経過するたびに勾配を更新します。これにより実質的な訓練バッチサイズは
hparams.batch_size.train
*hparams.accumulate_grad_batches
になります。
- この数値と同じstep数が経過するたびに勾配を更新します。これにより実質的な訓練バッチサイズは
dataset
:precomputed
を指定すると、事前にシリアライズされた記事を学習に用います。毎stepごとに行われる記事データのシリアライズ処理をスキップできるので、学習が高速化されます。
num_workers
:- DataLoaderのworker数(default:
4
)
- DataLoaderのworker数(default:
本ライブラリではコマンドラインパーサとしてhydraを用いているため、configで定義されている値は全てコマンドラインから書き換え可能になっています。
pipenv run tensorboard
host OS上のブラウザでlocalhost:${TENSORBOARD_PORT}
にアクセスするとログが確認できます。