1. 概要

1-1. 経緯

深層学習論文の再現実装をしたいと思っており、中でもVGGの論文はまぁまぁ簡単だと聞いて、VGGの再現実装に挑戦していたのですが、以下のような経緯で問題にぶつかりました。

ネットワークモデルを、 torchvision の提供する torchvision.models.vgg16 と完全に同じ状態にした。
※モデルの状態は、以下の[実装したネットワークモデル]を参照
torchvisionのvgg16と、自作したVGG16をそれぞれ学習させたところ、自作のVGG16のみ、学習が全然収束しなかった。
何回学習させても、「2.」の問題が解消せず。

[実装したネットワークモデル]

VGG

  VGG(
    (features): Sequential(
      (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (1): ReLU(inplace=True)
      (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (3): ReLU(inplace=True)
      (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
      (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (6): ReLU(inplace=True)
      (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (8): ReLU(inplace=True)
      (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
      (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (11): ReLU(inplace=True)
      (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (13): ReLU(inplace=True)
      (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (15): ReLU(inplace=True)
      (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
      (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (18): ReLU(inplace=True)
      (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (20): ReLU(inplace=True)
      (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (22): ReLU(inplace=True)
      (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
      (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (25): ReLU(inplace=True)
      (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (27): ReLU(inplace=True)
      (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (29): ReLU(inplace=True)
      (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    )
    (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
    (classifier): Sequential(
      (0): Linear(in_features=25088, out_features=4096, bias=True)
      (1): ReLU(inplace=True)
      (2): Dropout(p=0.5, inplace=False)
      (3): Linear(in_features=4096, out_features=4096, bias=True)
      (4): ReLU(inplace=True)
      (5): Dropout(p=0.5, inplace=False)
      (6): Linear(in_features=4096, out_features=1000, bias=True)
    )
  )

1-2. 解決策

自分の実装の間違っている箇所を調べるため、Githubに上がっている torchvision のソースコードを見に行った。

その結果、 モデルの重みパラメータの初期化 処理の有無が原因であったことが判明した。
重みパラメータの初期化処理を加えたところ、二つのモデルは似たような学習曲線を描いた。
（ただし、途中でColaboratoryの使用時間制限に引っかかって処理が停止したので、最後までは確認できず....）。

1-3. 結論

学習が収束しない原因は重みパラメータの初期値だった。
本で読んだことはあったが、実際に体験できたのは良い経験になった。

再現実装するなら、モデルだけじゃなくて重みパラメータの初期値も再現しないと、全然異なる学習結果になってしまう。

※当然ではあるが、その他にもハイパーパラメータなども再現しないとダメ

2. 詳細

2-1. 重みパラメータの初期化処理

torchvision側のソースと、自分が行った実装の違いはここだった。

PyTorch の vgg16 に実装されていた処理

# 重みパラメータを初期化している
        if init_weights:
            for m in self.modules():
                if isinstance(m, nn.Conv2d):
                    nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
                    if m.bias is not None:
                        nn.init.constant_(m.bias, 0)
                elif isinstance(m, nn.BatchNorm2d):
                    nn.init.constant_(m.weight, 1)
                    nn.init.constant_(m.bias, 0)
                elif isinstance(m, nn.Linear):
                    nn.init.normal_(m.weight, 0, 0.01)
                    nn.init.constant_(m.bias, 0)

該当箇所のソースコードをGithub上で確認したい場合は＊ココ＊をクリックして下さい

2-2. 解説

わかる範囲で上記ソースを解説します。
適宜、Githubのソースと比較しながら見ていただけるとわかりやすいと思います。

2-2-1. if init_weights:

最初の init_weights は、torchvision.models.vgg16(weights=None)にした場合にTrueになります。

weights引数に対して、None以外の値を与えた場合は、init_weights=Falseになっているようでした。

2-2-2. for m in self.modules:

self.modulesには、各層がmodule化されて入っています。
以下のような感じ。

self.modulesの中身

# from torch import nn が実行されているとする

nn.Sequentials(
    nn.Conv2d(...),
    nn.ReLU(...),
    nn.Conv2d(...),
    ...
)

なので、これがforループで回されると、

1ループ目がConv2dのmodule
2ループ目がReLUのmodule
...

という感じになります。

2-2-3. if isinstance(m, nn....):

isinstanceは、pythonの組み込み関数ですね。
mにはConv2dなどが入っています。

このmが、 Conv2d BatchNorm2d Linear のいずれかから作られたインスタンスだったときだけ if条件がTrue になります。

以降の 2-2-4、2-2-5、2-2-6 で取り上げる nn.init 以下の関数については、以下の記事で説明がなされていました。

この記事はわかりやすいですし、これらの関数を使うだけなら、この記事の説明で十分足りると思います。

2-2-4. nn.init.kaiming_normal_(m.weights, ...)

恐らく、一番大事なところです。

pytorchの公式ドキュメントがとても親切で、使い方まで含めて記載されているので、リンクを貼っておきます。ぺたり。

torch.nn.init.kaiming_normal_のドキュメント

kaiming_normal_ メソッドの概要

いわゆる Heの初期値 と呼ばれるものを使ってモデルパラメータの初期化を行ってくれるメソッドです。
normal_というpostfixがついていることからわかるように、正規分布に従う値でパラメータを初期化してくれます。

本メソッドの各引数については、上記 kaiming_normal_ のドキュメントと、pytorchの kaiming_normal_ソースを見て、確認しました。

第一引数

kaiming_normal_は、第一引数に与えた Tersor型のパラメータを初期化 してくれます。

引数 `a`

今回vgg16では、mode nonlinearityのみが指定されていたので、引数aについての話は省略。

引数 `mode`

デフォルト値の 'fan_in' か、もしくは 'fan_out' が選べるようです。

'fan_in'の場合は、一つ前の層の分散の大きさが維持されるように、モデルパラメータを初期化するようです。
対して'fan_out'を指定した場合は、一つ後ろの層の分散の大きさが維持されるようにモデルパラメータを初期化するようです。

ソースコードはこちら

TODO: なぜVGGでfan_outが指定されているのか、わかっていない

引数 `nonlinearity`

'relu' か 'leaky_relu' を指定した方が良いよとドキュメントには書かれています。

今回のvgg16ではReLUを使用したので、'relu'を指定しておくので良さそう。

nonlinearity引数として他に何が指定できるのか調べてみたところ、sigmoidやtanhなども指定できるようです。

nonlinearityに指定した各引数による動作の違いは、以下のソースを参照。

2-2-5. nn.init.constant_(m.bias, 0)

各層のバイアス(bias)の初期値を定数(constant)で初期化する。

nn.init.constant_(m.bias, 0)のように書いた場合は、全て0で初期化。

2-2-6. nn.init.normal_(m.weight, 0, 0.01)

各引数の意味合いは公式ドキュメントを見るとすぐにわかる。

torch.nn.init.normal_(tensor, mean=0.0, std=1.0)

と書かれているので、nn.init.normal_(m.weight, 0, 0.01)と記述した場合は、平均0、標準偏差0.01の正規分布からサンプリングされた値で、第一引数のm.weightを初期化してくれる。

3. 補足

今回の torchvision.models.vgg16 で行われている重みパラメータの初期化はVGGの論文に記載されているものとは異なります。

VGGの論文では、"3 CLASSIFICATION FRAMEWORK > 3.1 TRAINING" の第2パラグラフで重みパラメータの初期化について触れられており、それによると以下の手順で初期化が行われたようです。

まず、比較的層数の少ない設定Aのモデルに、ランダムな初期パラメータを与えて学習
Aよりも深い層を持つ設定のモデルを学習させる際は、最初の4つのConvolution層と、最後の3つの全結合層に対して、「1.」で得られたパラメータを適用
その他の中間層は、ランダムなパラメータで初期化
「1.」や「3.」において、ランダムなパラメータで初期化する部分には、平均0、分散0.01の正規分布からランダムにサンプリングした値を適用
（バイアスは0で初期化）

学習時間がかかってしょうがないので、今回は省略しました(；˘ω˘)ｽｯ…ｽﾔｧ…

4. 感想など

VGGの再現実装を行うには 2-2. に書いたことを全て理解する必要があったらしい。

今回のように 学習が収束しない 場合、モデルパラメータの初期値に問題があるかもしれない とわかったのは、大きな収穫だった。

これ以外にも、論文によっては、モデルのネットワークやハイパーパラメータの他に、前処理、最適化関数や損失関数なども工夫していたりするので、まだまだ学ぶことはたくさんありそうだ。

他の論文も再現実装していくぞ！( *˙︶˙*)و

5. 参考

Heの初期値についての原著論文(Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification)は、以下のリンクから飛べます。

また、VGG論文中で触れられている、重みパラメータのランダムな初期化方法（Glorot & Bengio (2010)）については以下で確認できます。

より理解を深めるため、どっちも軽く目を通しておきたいです。

その他、論文全体の日本語訳は以下のQiita記事にまとまっていました。
大変ありがたく活用させていただきました！