google-research の vision_transformer を動かすdockerイメージを作成した
なぜ：
google-research の vision_transformerを動かすのにjaxが必要だが、gpu対応のjaxバイナリが提供されておらず、ビルドが必要かつ、ビルドにものすごい時間がかかったため、vision_transformerを動かすためのdockerイメージを作成した。
下記でvision_transformerを実行できる
$ docker run --name vision_transformer --gpus all notfolder/vision_transformer --help
ウチのGTX2070だとメモリが8GBしかないので下記コマンドでやっとこさ動く。
docker run --name vision_transformer --gpus all -v /home/notfolder/work/:/work notfolder/vision_transformer --name ViT-B_16-cifar100_`date +%F_%H%M%S` --model ViT-B_16 --logdir /work/vit_logs --output /work/vit_output/ViT-B_16-cifar100_`date +%F_%H%M%S` --total_steps=10000 --warmup_steps 100 --dataset cifar100 --accum_steps=1 --batch 4 --batch_eval 128
このdockerイメージをビルドするためのリポジトリが下記。
直接作るとビルドのためのいろいろなモジュールが入って大きくなりすぎるため、
一旦jaxlibのビルド環境を作り、その上でいったんjaxlib-0.1.58-cp36-none-manylinux2010_x86_64.whlを作ってそれを使ってdockerコンテナをビルドする。
https://github.com/notfolder/vision_transformer-docker
pretrainデータを自動ダウンロードする機能付きentorypointも作りました。