mucunwuxian
commented
3 years ago リサーチ備忘📝
調査した結果につきまして、こちらに記載させて頂こうと思います。
「lucidrains/vit-pytorh」について
lucidrains/vit-pytorchについて、確認させていただきました。 通常のvitだけでなく、複数種類のvitも登録されており、素晴らしいリポジトリとなっています。
Licenseは、MITになります。
リポジトリの構成としては、以下のようになっています。
./
├── LICENSE
├── README.md
├── examples
│ └── cats_and_dogs.ipynb
├── images
│ ├── cait.png
│ ├── cross_vit.png
│ ├── cvt.png
│ ├── dino.png
│ ├── distill.png
│ ├── levit.png
│ ├── nest.png
│ ├── pit.png
│ ├── t2t.png
│ ├── twins_svt.png
│ └── vit.gif
├── setup.py
└── vit_pytorch
├── __init__.py
├── cait.py
├── cross_vit.py
├── cvt.py
├── deepvit.py
├── dino.py
├── distill.py
├── efficient.py
├── levit.py
├── local_vit.py
├── mpp.py
├── nest.py
├── pit.py
├── recorder.py
├── rvt.py
├── t2t.py
├── twins_svt.py
└── vit.pyimages フォルダに格納されている画像は、readmeの説明に用いられているものとなります。
vit_pytorch フォルダに格納されているモジュールは、各種vitネットワーク構造のdefや、その補助機能等になります。
example フォルダに格納されている cats_and_dogs.ipynb というファイルは、当該リポジトリのコードを用いてスクラッチで学習する方法をレクチャーしているものでした。
Kaggleにある Dogs vs. Cats Redux: Kernels Edition というデータセットにて、学習する例が実装してあります。
readmeを読んでいくと、pre-trainedモデルのリンク紹介があった為、覗いてみると以下の別リポジトリへと繋がりました。
↓
https://github.com/rwightman/pytorch-image-models
しかし、当該リポジトリとリンク先リポジトリとを読み込んでみると、どうもvitそのもののpretrainedパラメーターは無さそうでした。 resnet等の有名なネットワークを、特徴抽出として用いるvitの場合に、その特徴抽出部分について、pretrainedモデルを用いれるというニュアンスのようでした。 或いは、readmeに記載のある学習サンプルにおいては、torchvisionのpretrainedモデルを用いて、蒸留を行う手順が記載されていたりしました。 https://github.com/lucidrains/vit-pytorch#distillation
その為、一旦ではありますが、ailia上で動かすonnxの生成元として、このリポジトリは不十分である印象を受けました。 他のリポジトリについても、調査をしてみようと思います。
「jeonsworld/ViT-pytorch」について
jeonsworld/ViT-pytorchというリポジトリについても、確認させていただきました。 こちらも☆の数が600を超えるリポジトリで、中身も充実したものでした。
Licenseは、MITになります。
リポジトリの構成としては、以下のようになっています。
./
├── LICENSE
├── README.md
├── img
│ ├── figure1.png
│ ├── figure2.png
│ └── figure3.png
├── models
│ ├── configs.py
│ ├── modeling.py
│ └── modeling_resnet.py
├── requirements.txt
├── train.py
├── utils
│ ├── data_utils.py
│ ├── dist_util.py
│ └── scheduler.py
└── visualize_attention_map.ipynbimg フォルダに格納されている画像は、readmeの説明に用いられているものとなります。
models フォルダに格納されているモジュールは、各種vitネットワーク構造のdefや、そのcofig補助機能等になります。
utils フォルダに格納されているモジュールは、学習の際に使用するloaderやscheduler等のdefとなります。
visualize_attention_map.ipynb というファイルは、当該リポジトリのコードを用いて、学習済みモデルにてfeed forwardを行うnotebook形式のチュートリアルとなっています。
学習済みモデルのダウンロードについても、コードでの実施が行われるように実装がされています。
また、attention mapによる、注目領域の可視化も行ってくれます。
結果も見せ方も良いコードとなっているので、こちらの踏襲すれば良いかと思ったのですが、よくよく調べてみると、このリポジトリにおけるモデルの実装は、Vision Transformer ではなく、Compact Convolutional Transformer というもののようでした。
その為、ひょっとすると、踏襲させていただくのは微妙かもしれないと思いました。
尚、補足としまして、当該リポジトリを用いた学習については、事前学習済みモデルを用いて学習する方針が推奨されておりますが、逆にスクラッチでの学習を行おうとした場合には、一部コードの改変を要するようでありました。 https://github.com/jeonsworld/ViT-pytorch/issues/7
その為、スクラッチで学習をする場合には、lucidrains/vit-pytorchの方が適しているかもしれません。
「google-research/vision_transformer」について
lucidrains/vit-pytorchのリポジトリを見ていると、officialリポジトリが以下である、という記載がありました。 https://github.com/google-research/vision_transformer
このリポジトリは、google-researchによるJAXの実装になります。 notebookによる実装もあり、Colab上にてTPUを使いながら、学習から推論までを一通り実行することができました。 可視化についても、論文の方針と合致しています。 https://colab.research.google.com/github/google-research/vision_transformer/blob/master/vit_jax.ipynb
構造は以下のようになっています。
./
├── CONTRIBUTING.md
├── LICENSE
├── README.md
├── figure1.png
├── vit_jax
│ ├── __init__.py
│ ├── checkpoint.py
│ ├── checkpoint_test.py
│ ├── configs.py
│ ├── flags.py
│ ├── hyper.py
│ ├── input_pipeline.py
│ ├── logging.py
│ ├── models.py
│ ├── models_resnet.py
│ ├── models_test.py
│ ├── momentum_clip.py
│ ├── requirements.txt
│ ├── train.py
│ └── train_test.py
└── vit_jax.ipynbmodels.py に、vitネットワーク構造の大枠が定義されています。
vit_jax.ipynb が、notebookによるスクリプトになります。
Licenseは、Apache-2.0になります。
このリポジトリのモデルを、onnxエクスポートできれば、ailiaへの搭載が可能と考えたのですが、jaxからonnxのエクスポートについて、情報が見つかりませんでした。 当該リポジトリ上のreadmeやissueにonnxというキーワードが無く、「jax onnx」などで検索をしても解決策は見当たらない形です。
しかし、色々調べてみると、JAX版のgoogle researchによるvitを、pytorchへとconvertしてくれているリポジトリを見つけました。 その為、そのリポジトリにて、pytorch版の学習済みモデルを作った後、onnxへのconvertを実施すれば、originalのvit構造を、ailiaに移植できると考えました。
「lukemelas / PyTorch-Pretrained-ViT」について
JAX版のgoogle researchによるvitを、pytorchへとconvertしてくれているリポジトリが以下になります。 https://github.com/lukemelas/PyTorch-Pretrained-ViT/issues?q=is%3Aissue+
尚、リポジトリに格納されているモジュール外に、Colabのスクリプトモジュールもあり、そのリンクはreadmeに記載されています。 https://colab.research.google.com/drive/1muZ4QFgVfwALgqmrfOkp7trAvqDemckO?usp=sharing
リポジトリの構造は以下の様になっています。
./
├── README.md
├── examples
│ ├── imagenet
│ │ ├── README.md
│ │ ├── data
│ │ │ └── README.md
│ │ └── main.py
│ └── simple
│ ├── example.ipynb
│ ├── imagenet-21k-labels.py
│ ├── imagenet.synset.obtain_synset_list
│ ├── img.jpg
│ ├── img2.jpg
│ ├── labels_map.txt
│ └── labels_map_21k.txt
├── jax_to_pytorch
│ ├── README.md
│ ├── convert.py
│ ├── explore-conversion-21k.ipynb
│ ├── explore-conversion.ipynb
│ ├── jax_weights
│ │ └── download.sh
│ └── weights
├── pytorch_pretrained_vit
│ ├── __init__.py
│ ├── configs.py
│ ├── model.py
│ ├── transformer.py
│ └── utils.py
└── setup.py作られて7ヶ月のリポジトリであるからか、実施内容の性質上か、Licenseについての記載はありませんでした。
convertの方法は非常にシンプルで、jaxで保存されたモデルを、pytorchのstate_dictへと変換するもののようでした。 構造が似ているために、微調整をすれば変換ができるようです。 https://github.com/lukemelas/PyTorch-Pretrained-ViT/blob/master/jax_to_pytorch/convert.py
この変換を経て、pytorchモデルからonnxエクスポートを行えば、googleresearchオリジナルのvitがailiaへと搭載できそうです。
https://github.com/lucidrains/vit-pytorch