Google Colab で RWKV のゼロからの学習を試す

「Google Colab」で「RWKV」のゼロからの学習 (開始するところまで) を試しました。

1. RWKV のゼロからの学習を試す

Google Colabでの実行手順は、次のとおりです

(1) RWKV-LMのリポジトリ経由でインストール。

!git clone https://github.com/BlinkDL/RWKV-LM
%cd RWKV-LM/RWKV-v4
%pip install transformers pytorch-lightning==1.9 deepspeed wandb ninja

(2) RWKV-LMフォルダ直下にdataフォルダを作成し、データセットを配置。
今回は、「https://data.deepai.org/enwik8.zip」をダウンロードして解凍して配置します。「enwiki8」は、2006年3月3日の英語版ウィキペディアXMLダンプの最初の100Mのデータセットです。

(3) 「RWKV-v4」の「train.py」の設定を確認。
今回は、そのまま実行します。

datafile = "../data/enwik8" # your data
datafile_encoding = 'utf-8' # 'utf-8' / 'utf-16le' / 'numpy' (for fine-tuning pile models) / 'binidx' (the Megatron-LM 'binidx' format)

EPOCH_BEGIN = 0 # begins with miniEpoch = EPOCH_BEGIN
LOAD_MODEL = False # shall we load the #EPOCH_BEGIN model and continue the training from it?

n_layer = 6
n_embd = 512
ctx_len = 1024 # increase T_MAX in src/model.py if your ctx_len is longer

model_type = 'RWKV' # 'RWKV' or 'RWKV-ffnPre' (sometimes better)

(4) 学習の実行。

# 学習の実行
!python train.py

デフォルトでは、「trainX.pth」という名前で、10エポック毎にモデルが保存されます。

Unityではじめる ROS・人工知能 ロボットプログラミング実践入門