AI

kernel_initializerって学習の収束に大事かも

例によってKerasからPyTorchへのお引越しでプチハマったことがあったのでまとめておきます。

Kerasでは学習が簡単に収束したのに、同じようにプログラムを組んだつもりでもPyTorchでは簡単には収束しない(あるいは同じ性能が出ない)ことってありますよね。今回グッディーは、KerasではoptimizerにRMSprop()を指定していい感じで収束していたのが、PyTorchでも同じようにRMSprop()を指定してもまったく収束しない、という場面に遭遇しました。

Adam()に変えたら無事収束したのでまあいいかと思っていましたが、やはり気になったのでもう少し違いを深堀りしてみたというのが事のいきさつです。

では実際のコードを見てみましょう…

Kerasではこんな感じ。

model = Sequential() 
model.add(Conv1D(filters, kernel_size, strides, padding='valid', input_shape, dilation_rate=1, activation=None, use_bias=True, kernel_initializer='glorot_uniform', bias_initializer='zeros', kernel_regularizer=None, bias_regularizer=None, activity_regularizer=None, kernel_constraint=None, bias_constraint=None))
model.add(BatchNormalization())
model.add(Bidirectional(LSTM(n_hidden, dropout)))
model.add(Dense(lnum, kernel_initializer='truncated_normal'))
model.add(Dropout(dropout))
model.add(Activation('softmax'))

model.compile(loss='mae', optimizer=RMSprop(), metrics)

イニシャルのPyTorchでの実装はこんな感じ。(修正部分は torch.viewも気をつけて使おう ! を参照)

class LSTM(nn.Module): 
    def init(self, input_size, outch_size, hidden_layer_size, output_size): 
        super().init() 
        self.conv = nn.Conv1d(in_channel, out_channel, kernel_size, stride=1, padding=0)
        self.bn = nn.BatchNorm1d(outch_size, momentum, eps) 
        self.lstm = nn.LSTM(outch_size, hidden_layer_size, batch_first=True, dropout, bidirectional=True)
        self.linear = nn.Linear(hidden_layer_size * 2, output_size) 
        self.do = nn.Dropout(dropout) 
        self.sm = nn.Softmax(dim=1) 

    def forward(self, input_sq): 
        cf_sq = self.conv(input_sq) 
        ncf_sq = self.bn(cf_sq) 
        batch_size, feature_size, sq_len = ncf_sq.shape 
        lstm_out, _ = self.lstm(ncf_sq.view(batch_size, sq_len, feature_size)) 
        lstm_out, _ = self.lstm(ncf_sq.permute(0, 2, 1)) 
        out = self.linear(lstm_out[:, -1, :]) 
        out = self.do(out) 

        predictions = self.sm(out) 
        return predictions 

# main() 
    model = LSTM(input_size, hidden_layer_size, output_size) 
    loss_function = nn.MSELoss() 
    optimizer = torch.optim.RMSprop(model.parameters()) 

違いはというと…

KerasではConv1D()やDense()の引数でkernel/bias_initializerやkernel/bias_regularizerを指定できます。一方PyTorchではそのような指定はありません。ここでは以下の2点に注目したいと思います。

  • Conv1D()のkernel_initializer=’glorot_uniform’
  • Dense()のkernel_initializer=’truncated_normal’

ちなみに、xxx_normalは正規分布系、xxx_uniformは一様分布系、trancatedは分布の端を切断したものになります。たとえばこのあたり (Keras Documentation) が参考になります。

ではPyTorchではどのようなInitializerが用いられているのでしょうか ? CrossEntropyLoss – Expected object of type torch.LongTensor に以下のようなコメントがあります。

Also, PyTorch initializes the conv and linear weights with kaiming_uniform and the bias with a uniform distribution using fan_in by default.

また下の方にこのようなコメントもありますね。

Yes, he_uniform would correspond to kaiming_uniform in PyTorch and I think I was referring to the default Keras initialization based on their docs, which is glorot_uniform (so xavier_uniform in PyTorch).

ということで、PyTorchではデフォルトでkaiming_uniformすなわちhe_uniformのkernel_initializerが適用されるとのこと。

glorot_uniformの適用

そして記事にあるとおり、xavier_uniformをPyTorchのConv1d()に適用してあげると、KerasのConv1D()のkernel_initializer=’glorot_uniform’になるので、bias_initializer=’zeros’の部分も含めて提案されているサンプルコードをそのまま適用すればよさそうです。

def weight_init_glorot_uniform(m): 
    if isinstance(m, nn.Conv1d): 
        nn.init.xavier_uniform_(m.weight, gain=g1) 
        nn.init.zeros_(m.bias) 

model.apply(weight_init_glorot_uniform)

truncated_normalの適用

ではtruncated_normalの方はというと、Implementing Truncated Normal Initializer での議論がとても役に立ちます。グッディーのケースでは厳密にやる必要はなかったので、中ほどで提案されている近似手法を適用してみました。

def weight_init_truncated_normal(m): 
    if isinstance(m, nn.Linear): 
        size = m.weight.size() 
        z = torch.fmod(torch.randn(size), g2) 
        m.weight.data = z 

model.apply(weight_init_truncated_normal)

結果

あとは双方の関数定義中にあるゲイン (g1, g2) を決めてやる必要があります。TORCH.NN.INIT を読むと、

  • nn.init.calculate_gain(‘linear’) = 1
  • nn.init.calculate_gain(‘relu’) = \sqrt{2}

となっているのと、上の近似式の記事ではg2が2になっていたので、1, \sqrt{2}, 2の3種類で実験してみました。結果は…

  • g1=\sqrt{2}, g2=1 : 収束
  • g1=\sqrt{2}, g2=\sqrt{2} : 急速に収束
  • g1=2, g2=1 : 収束
  • g1=2, g2=2 : 緩やかに収束
  • それ以外 : 収束せず

結構lossやaccuracyの収束曲線の形状は違っていました。しかし収束値はほぼ同等という結果でした。

さいごに

ということで、PyTorchのConv1d()とLinear()にそれぞれglorot_uniformとtruncated_normalのkernel_initializerを適用し、ゲインを変えて収束状態を比較してみました。

モデルやオプティマイザーにも大きく依存するのだと思いますが、グッディーのケースでは収束するかしないかにかなりセンシティブな印象を受けました。一方収束する場合はその収束速度にばらつきはあるものの、収束値はほぼ同等でした。

ということで、作成したモデルの学習が収束しないとき、kernel_initializerを変えてみることは重要と言えます。興味がありましたら、あなたのモデルでも実験してみてください !


   
関連記事
  • PyTorchで学習済みモデルの中間層出力の取得
  • Ubuntu20.04+CUDA11.2+PyTorch1.8.1+cu111がRTX3090のGPUで動作した !
  • The NVIDIA driver on your system is too oldって !!!
  • Tensorflow 1と2のModel混在のときに目からうろこだった話
  • Pytorch Tutorial その1 – Tensor
  • PyTorchで突然malloc(): invalid next size (unsorted)が出たときの対処

    コメントを残す

    *

    CAPTCHA