EfficientNetの事前学習モデルをPytorchで使う

#何の記事？
EfficientNetの事前学習モデルをKerasを用いて動かす方法は、こちらで解説されていますが、今回、Pytorchでも動かす方法を見つけたので、共有します。

#EfficientNetとは？
2019年5月にGoogle Brainから発表されたモデルです。広さ・深さ・解像度を効率よくスケールアップすることにより、少ないパラメータ数で高い精度を実現しています。使用できる単体のモデルの中では、最強と言ってよいでしょう。

>Tan, Mingxing, and Quoc V. Le. "EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks." arXiv preprint arXiv:1905.11946 (2019).

#Pytorch用のpretrained model
パラメータは、こちらのGitHub内に保存されています。
https://github.com/lukemelas/EfficientNet-PyTorch

以下のように書くと、そのままモデルにインポートできます。"efficientnet-b7"の数字の部分は、0～7まで変更可能ですが、上のグラフからわかるように、"7"を使うのが最も精度が高いです。

!pip install efficientnet_pytorch 　#Google colabだとこれで動きます
from efficientnet_pytorch import EfficientNet

model_ft = EfficientNet.from_pretrained('efficientnet-b7') 　#Pretrained_modelのインポート

事前学習にはImageNetを用いているので、全結合層は1000クラスです。以下のようにクラス数を変更すれば転移学習に利用することができます。

num_ftrs = model_ft._fc.in_features　　#全結合層の名前は"_fc"となっています
model_ft._fc = nn.Linear(num_ftrs, 2)

###Pytorchを用いた転移学習
転移学習のスクリプトについて知りたい方は、こちらのサイトを参考にして下さい。
Pytorch tutorial： https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html