ELECTRAのpre-trainingを実行

前準備

Tensorflow(GPU版)をpipでインストール。pipのバージョン違いで混乱することがあるため、私はpython -m pipでインストールした。バージョンは1.15にすること。

python -m pip install tensorflow-gpu==1.15

gitでELECTRAのリポジトリをクローン

git clone https://github.com/google-research/electra.git

vocab.txtをダウンロード

cd electra
wget https://storage.googleapis.com/electra-data/vocab.txt
cp vocab.txt /home/ubuntu/disk/openwebtext

OpenWebTextCorpus(12G)をGoogle driveからダウンロードし、以下で解凍

tar xf openwebtext.tar.xz

この時点で、フォルダ構成は以下のようになっている。私の場合は、プログラムを /home/ubuntu/script/electra 以下におき、データは /home/ubuntu/disk/openwebtext においた

/home/ubuntu/script/electra 以下のディレクトリ構成

electra ├── build_openwebtext_pretraining_dataset.py ├── build_pretraining_dataset.py ├── configure_finetuning.py ├── configure_pretraining.py ├── CONTRIBUTING.md ├── finetune │   ├── classification │   │   ├── classification_metrics.py │   │   └── classification_tasks.py │   ├── feature_spec.py │   ├── init.py │   ├── preprocessing.py │   ├── qa │   │   ├── mrqa_official_eval.py │   │   ├── qa_metrics.py │   │   ├── qa_tasks.py │   │   ├── squad_official_eval.py │   │   └── squad_official_eval_v1.py │   ├── scorer.py │   ├── tagging │   │   ├── tagging_metrics.py │   │   ├── tagging_tasks.py │   │   └── tagging_utils.py │   ├── task_builder.py │   └── task.py ├── flops_computation.py ├── LICENSE ├── model │   ├── init.py │   ├── modeling.py │   ├── optimization.py │   ├── pycache │   │   ├── init.cpython-36.pyc │   │   ├── modeling.cpython-36.pyc │   │   ├── optimization.cpython-36.pyc │   │   └── tokenization.cpython-36.pyc │   └── tokenization.py ├── pretrain │   ├── init.py │   ├── pretrain_data.py │   ├── pretrain_helpers.py │   └── pycache │   ├── init.cpython-36.pyc │   ├── pretrain_data.cpython-36.pyc │   └── pretrain_helpers.cpython-36.pyc ├── pycache │   ├── build_pretraining_dataset.cpython-36.pyc │   └── configure_pretraining.cpython-36.pyc ├── README.md ├── run_finetuning.py ├── run_pretraining.py ├── util │   ├── init.py │   ├── pycache │   │   ├── init.cpython-36.pyc │   │   ├── training_utils.cpython-36.pyc │   │   └── utils.cpython-36.pyc │   ├── training_utils.py │   └── utils.py └── vocab.txt

/home/ubuntu/data/openwebtext 以下のディレクトリ構成

openwebtext ├── openwebtext │   ├── urlsf_subset00-1_data.xz │   ├── urlsf_subset00-2_data.xz │   ├── urlsf_subset00-3_data.xz │   ├── urlsf_subset00-4_data.xz │   ├── urlsf_subset00-5_data.xz │   └── (以下省略) └── vocab.txt

プログラムの実行

まず、学習用のデータを整形するため、以下のスクリプトを実行

python build_openwebtext_pretraining_dataset.py --data-dir /home/ubuntu/data/openwebtext --num-processes 5

実行後、/home/ubuntu/data/openwebtext 以下にpretrain_tfrecordsという名前のフォルダとその下にpretrain_data.tfrecord-0-of-1000のようなファイルが1000個生成される。また、tmpフォルダには、テキスト化されたデータが入っているので、ここで内容を確認することもできる。

ここまでくれば、いよいよpre-trainingが実行できる。

python run_pretraining.py —data-dir /home/ubuntu/data/openwebtext —model-name electra_small_owt

上のコマンドを実行すると学習が始まる。GithubのReadmeによると、openwebtextの学習には、v100を使って4日ほどかかるらしいので、最後まで学習するのは多くの人にとって困難だと思います。

ちなみに筆者は、g4dn16xlargeで5時間くらいまわして、1%ほど学習が進んだことを確認して止めました。最後に、学習後のディレクトリ構成をスクショで共有させていただき、本記事を締めさせていただきたいと思います。

次の記事では、build_openwebtext_pretraining_dataset.py によって生成されたtmpフォルダ内に存在する、テキストデータを使って、自作テキストから事前学習を行う方法について書いていきたいと思います。

学習後は新たにmodelsフォルダが生成