世界モデルの試み（I-JEPA）を、非力なＧＰＵ（T4）で動かしてみます。
背景
ルカン氏が、世界モデルのプロジェクトを推進することを公表しました。［※１］［※２］［※３］［※４］
その一歩として、まず画像を使った「予測」部分の実装（I-JEPA ）を試みています：［※５］
- ◯
- 公式（世界モデルの試み：I-JEPA）
- ・
- https://ai.facebook.com/blog/yann-lecun-ai-model-i-jepa/
- ◯
- 論文（世界モデルの試み：I-JEPA）［※６］［※７］
- ・
- https://arxiv.org/2301.08243
- ◯
- 実装（世界モデルの試み：I-JEPA）
- ・
- https://github.com/facebookresearch/ijepa
- ※１
- 汎用人工知能（ＡＧＩ）は、ヒトのように学習〜推論できることを目指しています（貧しい刺激からの学習、飛躍した／内省する推論）。いっぽうヒトの学習〜行動については、予測符号化や自由エネルギー原理（変分自由エネルギーを最小にするよう、主体は学習〜行動する）が説明する、という考え方があったりしますーーこの主張によれば、主体（ヒト）は自律的に、予測によって環境のモデルを構築するわけですがーーそして世界モデルも、主体（ＡＩ）が自律的に、予測によって環境のモデルを構築することを目指しています：
- ・
- https://www.riken.jp/press/2022/20220114_3/
- ※２
- 世界モデルは、モデルベースの強化学習でありながら、そのモデル（環境のモデル）を主体自身が構築する、という点に特徴がありますーーいっぱんに強化学習では、主体は環境の状態を観察し〜報酬を得て、行動を調整しますよねーーこのうちモデルフリーの強化学習は、環境についての知識がないので、行動にムダやヌケが多く、また訓練したＮＮモデルも（その主体と環境に特化しているので）汎用性がありません。いっぽうモデルベースの強化学習は、環境のモデルを持ち、その状態や報酬を予測しつつ行動できるので、ムダやヌケを最小に抑えることができます（それぞれの環境に特化した訓練も、環境のモデルを使って行えるので、より効率的になります）。とはいえモデルベースの場合、環境のモデルをどう構築するかが問題になります（自律的な主体なら、モデルを作るのは、その主体以外にないのでーーあらかじめモデルを作っておくというなら、それを行うべつの主体が必要になり、無限後退に陥りますし）ーー世界モデルでは、この環境のモデルの構築を、自己教師アリ学習などで＜主体と環境の相互作用＞を＜もつれのない低次元の埋め込み＞にし〜さらに＜予測＞を組み込むことで実現します（映像による処理ではVAE + MDN-RNN など）ーー環境のモデルは抽象化されているので、主体は、より単純なやり方で行動と報酬から学習することができます（進化計算のESやGAなど）。さらにその訓練の結果を環境に再適用し、それを元に環境のモデルを再構築するーーこのサイクルを繰り返すことで、環境に適応していくことができるかもしれません（最初に環境のモデルを構築するサンプルは、主体のランダムな行動による環境との相互作用になりますが、比較的単純な環境なら、この１回の訓練でも十分な効果を上げることは分かっています）：
- ・
- https://arxiv.org/abs/1803.10122
- ※３
- いまの大規模言語モデル（自己教師アリ学習で、分布仮説から次の単語を予測するモデル）は、大量のデータから学習させる必要があり、飛躍した／内省する推論もできないので、これ単体では、ヒトのような知的なふるまいは望めませんーーただ、このしくみをオフラインの強化学習に取り入れると、トランスフォーマのもつ＜系列の離れた位置にある相関＞を利用できるので、模倣学習の性能を飛躍的に高めることは分かっていますーーさらにこれを、メタ強化学習に応用した試みもあるようです：
- ・
- https://horomary.hatenablog.com/entry/2022/11/28/212738
- ・
- https://openreview.net/forum?id=hy0a5MMPUv
- ※４
- ルカン氏が考える世界モデルは、次に概要がまとめられていますーーすべてを微分可能なパーツで構成することに主眼をおいていますが、あくまで構想のレベルです（ここからすべてを設計〜実装できるようなものではありません）：
- ・
- https://ai.facebook.com/blog/yann-lecun-advances-in-ai-research/
- ※５
- この実装については、次が構成と手順を分かりやすく解説していますーー予測器を訓練し、画像の大きい部分から〜そのなかの小さい部分を予測します。ただし予測の対象に、それぞれの埋め込み（ベクトル化）を使っている点に特徴があり、ここがキモだったりします：
- ・
- https://www.youtube.com/watch?v=6bJIkfi8H-E
- ※６
- 埋め込み（ベクトル化）から予測するので、意味のより抽象的なレベル（深いセマンティクス）で訓練することができ、学習の効率を上げることができるようです（＝ヒトのような＜貧しい刺激からの学習＞に近づく）。反面、意味のより具体的なレベル（表面的な画像の再現）では、予測した箇所がズレるなど、詳細さは犠牲になりますが。
- ※７
- ほかの自己教師アリ学習モデルとの比較は、エネルギーベースモデル（非線形・非平衡の物理から流用した、最小の自由エネルギーを探索するモデル）で分類し、それらの利点／欠点を挙げていますーーこのモデル（I-JEPA）については、不変ベースや生成ベースのモデルより、セマンティクスが深く・汎用性も高い、という利点があるようです（その代わり、生成ベースより、詳細さでは劣る）。
概要
世界モデルもＲＮＮ系からトランスフォーマに移行し、これまでのような非力なＧＰＵでは実用にならない状況になっていますーーそしてこのプロジェクト（I-JEPA）の実装も、多くのリソースを要求します。［※１］
というわけでここではーーその非力なＧＰＵ（T4）でどこまで動かすことができるか、試してみます……
- ※１
- もちろん系列を並列・高速にあつかえるトランスフォーマだからこそ、ＲＮＮ系では実現できなかったことが可能になるんですけどね。
手順：導入
以下、導入の手順です。
データ用／ログ用のフォルダを作り：
$ cd ${dir_exp}
$ mkdir dat log
$ mkdir -p dat/imagenet_full_size/061417/train
$ mkdir -p log/vith14.224-bs.2048-ep.300
データ用のフォルダに、訓練用のデータ（ImageNet）を展開します：［※１］
$ cd ${dir_exp}/dat/imagenet_full_size/061417/train
$ cp -p /var/tmp/ILSVRC2012_train.tar.gz .
$ gzip -d ILSVRC2012_train.tar.gz
$ tar -xf ILSVRC2012_train.tar
$ rm ILSVRC2012_train.tar
ソースを取得し：
$ cd ${dir_exp}
$ git clone --depth=1 https://github.com/facebookresearch/ijepa.git
設定を変更します：
- ◯
- ${dir_exp}/ijepa/configs/in1k_vith14_ep300.yaml
root_path: $replace_this_with_absolute_path_to_your_datasets_directory folder: $replace_this_with_path_for_experiment_logs/vith14.224-bs.2048-ep.300/ model_name: vit_huge use_bfloat16: true batch_size: 128 num_workers: 10
root_path: ${dir_exp}/dat/
folder: ${dir_exp}/log/vith14.224-bs.2048-ep.300/
model_name: vit_large # ＧＰＵのメモリに依存
use_bfloat16: false # デバイスのサポートに依存
batch_size: 4 # メモリに依存
num_workers: 8
一部ソースの変更も必要でした：
- ◯
- ${dir_exp}/ijepa/src/train.py
with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=use_bfloat16):
with torch.cuda.amp.autocast(dtype=torch.float16, enabled=use_bfloat16):
なおビジョントランスフォーマ（ViT ）の定義は、次にあります：
- ◯
- ${dir_exp}/ijepa/src/models/vision_transformer.py
...
VIT_EMBED_DIMS = {
'vit_tiny': 192,
'vit_small': 384,
'vit_base': 768,
'vit_large': 1024,
'vit_huge': 1280,
'vit_giant': 1408,
}
...
とりえずコンテナの起動設定は最低限のものを記述し：
- ◯
- ${dir_doc}/Dockerfile
FROM ubuntu:22.04 RUN apt-get -y update RUN apt-get -y upgrade
コンテナを起動します（コンテナの変更は、適宜反映します）：
# 作成
$ docker build --no-cache -t ubuntu:vir124 ${dir_doc}
# 起動
$ docker run -it --rm --gpus all -v /opt:/opt --name cnt124 ubuntu:vir124
# 反映
$ docker commit cnt124 ubuntu:vir124
コンテナで、依存するライブラリ群を設置します：
$ apt install python3 $ apt install pip $ pip install torch==2.0 $ pip install torchvision $ pip install pyyaml $ #pip install numpy $ pip install opencv-python $ pip install submitit
- ※１
- すべてを展開すると非常に大きなサイズになるので（またいずれにしても実用にはならないので）、ここでは最初の２コほどのフォルダを展開していますーーImageNetのデータがなければ、適当なJPEG画像群を入れておいても、動作はするはずです。
手順：実行
次を実行し、訓練を始めます：
$ cd ${dir_exp}/ijepa
$ python3 main.py \
--fname configs/in1k_vith14_ep300.yaml \
--devices cuda:0
INFO:root:called-params configs/in1k_vith14_ep300.yaml
INFO:root:loaded params...
{ 'data': { 'batch_size': 4,
'color_jitter_strength': 0.0,
'crop_scale': [0.3, 1.0],
'crop_size': 224,
'image_folder': 'imagenet_full_size/061417/',
'num_workers': 8,
'pin_mem': True,
'root_path': '${dir_exp}/dat/',
'use_color_distortion': False,
'use_gaussian_blur': False,
'use_horizontal_flip': False},
'logging': { 'folder': '${dir_exp}/log/vith14.224-bs.2048-ep.300/',
'write_tag': 'jepa'},
'mask': { 'allow_overlap': False,
'aspect_ratio': [0.75, 1.5],
'enc_mask_scale': [0.85, 1.0],
'min_keep': 10,
'num_enc_masks': 1,
'num_pred_masks': 4,
'patch_size': 14,
'pred_mask_scale': [0.15, 0.2]},
'meta': { 'copy_data': False,
'load_checkpoint': False,
'model_name': 'vit_large',
'pred_depth': 12,
'pred_emb_dim': 384,
'read_checkpoint': None,
'use_bfloat16': False},
'optimization': { 'ema': [0.996, 1.0],
'epochs': 300,
'final_lr': 1e-06,
'final_weight_decay': 0.4,
'ipe_scale': 1.0,
'lr': 0.001,
'start_lr': 0.0002,
'warmup': 40,
'weight_decay': 0.04}}
INFO:torch.distributed.distributed_c10d:Added key: store_based_barrier_key:1 to store for rank: 0
INFO:torch.distributed.distributed_c10d:Rank 0: Completed store-based barrier for key:store_based_barrier_key:1 with 1 nodes.
INFO:root:Running... (rank: 0/1)
INFO:root:Initialized (rank/world-size) 0/1
INFO:root:VisionTransformer(
(patch_embed): PatchEmbed(
(proj): Conv2d(3, 1024, kernel_size=(14, 14), stride=(14, 14))
)
(blocks): ModuleList(
(0-23): 24 x Block(
(norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)
(attn): Attention(
(qkv): Linear(in_features=1024, out_features=3072, bias=True)
(attn_drop): Dropout(p=0.0, inplace=False)
(proj): Linear(in_features=1024, out_features=1024, bias=True)
(proj_drop): Dropout(p=0.0, inplace=False)
)
(drop_path): Identity()
(norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)
(mlp): MLP(
(fc1): Linear(in_features=1024, out_features=4096, bias=True)
(act): GELU(approximate='none')
(fc2): Linear(in_features=4096, out_features=1024, bias=True)
(drop): Dropout(p=0.0, inplace=False)
)
)
)
(norm): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)
)
INFO:root:making imagenet data transforms
INFO:root:data-path ${dir_exp}/dat/imagenet_full_size/061417/train/
INFO:root:Initialized ImageNet
INFO:root:ImageNet dataset created
INFO:root:ImageNet unsupervised data loader created
INFO:root:Using AdamW
INFO:root:Epoch 1
INFO:root:[1, 0] loss: 0.472 masks: 98.0 42.0 [wd: 4.00e-02] [lr: 2.00e-04] [mem: 1.07e+04] (2643.6 ms)
INFO:root:[1, 0] grad_stats: [4.57e-02 1.50e-02] (1.50e-02, 5.46e-02)
INFO:torch.nn.parallel.distributed:Reducer buckets have been rebuilt in this iteration.
INFO:torch.nn.parallel.distributed:Reducer buckets have been rebuilt in this iteration.
INFO:root:[1, 10] loss: 0.297 masks: 94.1 42.5 [wd: 4.00e-02] [lr: 2.00e-04] [mem: 1.09e+04] (910.5 ms)
INFO:root:[1, 10] grad_stats: [4.97e-04 8.87e-05] (8.87e-05, 7.68e-04)
INFO:root:[1, 20] loss: 0.231 masks: 94.5 43.1 [wd: 4.00e-02] [lr: 2.01e-04] [mem: 1.09e+04] (831.2 ms)
INFO:root:[1, 20] grad_stats: [8.38e-04 7.14e-05] (7.14e-05, 1.11e-03)
...
INFO:root:[1, 630] loss: 0.269 masks: 93.4 44.1 [wd: 4.00e-02] [lr: 2.19e-04] [mem: 1.09e+04] (761.2 ms)
INFO:root:[1, 630] grad_stats: [2.94e-05 2.34e-05] (2.33e-05, 8.50e-05)
INFO:root:[1, 640] loss: 0.270 masks: 93.4 44.1 [wd: 4.00e-02] [lr: 2.20e-04] [mem: 1.09e+04] (761.1 ms)
INFO:root:[1, 640] grad_stats: [1.30e-05 7.64e-06] (7.31e-06, 2.76e-05)
INFO:root:avg. loss 0.270
INFO:root:Epoch 2
INFO:root:[2, 0] loss: 0.268 masks: 102.0 42.0 [wd: 4.00e-02] [lr: 2.20e-04] [mem: 1.09e+04] (824.6 ms)
INFO:root:[2, 0] grad_stats: [1.76e-05 1.05e-05] (9.85e-06, 3.23e-05)
INFO:root:[2, 10] loss: 0.325 masks: 94.7 44.1 [wd: 4.00e-02] [lr: 2.20e-04] [mem: 1.09e+04] (749.0 ms)
INFO:root:[2, 10] grad_stats: [1.77e-05 1.33e-05] (1.28e-05, 5.03e-05)
INFO:root:[2, 20] loss: 0.324 masks: 95.0 43.7 [wd: 4.00e-02] [lr: 2.21e-04] [mem: 1.09e+04] (755.0 ms)
INFO:root:[2, 20] grad_stats: [8.32e-05 1.08e-05] (9.92e-06, 8.32e-05)
...
このときのＧＰＵの使用状況は：
$ nvidia-smi --query-gpu=timestamp,name,utilization.gpu,utilization.memory,memory.used,memory.free,memory.used --format=csv -l 1 ... 2023/06/17 12:29:24.274, Tesla T4, 100 %, 49 %, 14225 MiB, 741 MiB, 14225 MiB 2023/06/17 12:29:25.277, Tesla T4, 95 %, 97 %, 14225 MiB, 741 MiB, 14225 MiB 2023/06/17 12:29:26.279, Tesla T4, 100 %, 35 %, 14225 MiB, 741 MiB, 14225 MiB 2023/06/17 12:29:27.281, Tesla T4, 100 %, 80 %, 14225 MiB, 741 MiB, 14225 MiB ...