ShakedropのKeras実装

はじめに

AlgoAgeの林です
Shakedrop周りの論文をまとめたスライドと、Kerasの実装を掲載します。

shakedropについて

スライド（DLHacksで発表したもの）

実装

参考

・https://github.com/jonnedtc/Shake-Shake-Keras
・https://github.com/owruby/shake-drop_pytorch

shakedrop実装

shakedrop_model.py

import keras
from keras import Input
from keras import backend as K
from tensorflow import distributions as tfd

class Shakedrop(layers.Layer): # カスタムレイヤーを定義

    def __init__(self, num_of_unit, num_of_layers, **kwargs):
        super(Shakedrop, self).__init__(**kwargs)
        self.num_of_unit = num_of_unit # 何番目のresblockか
        self.num_of_layers = num_of_layers # モデル全体の層の数

    def build(self, input_shape):
        super(Shakedrop, self).build(input_shape)

    def call(self, x):
        batch_size = K.shape(x)[0]
        alpha = K.random_uniform((batch_size, 1, 1, 1),  minval=-1.0)
        beta = K.random_uniform((batch_size, 1, 1, 1))
        p = 1 - (self.num_of_unit / (2 * self.num_of_layers)) # 出力に近い方がshakeされやすくなる
        bernoulli = tfd.Bernoulli(probs=p).prob(1)

        def x_shake():
            # stop_gradientを使ってfowardとbackwardで切り替える
            return (1 - bernoulli) * (beta * x + K.stop_gradient((alpha - beta) * x))

        def x_even():
             # pがそのままbの期待値になる
            return p * x

        # 学習時はx_shakeでtestの時はx_even
        return K.in_train_phase(x_shake, x_even)

    def compute_output_shape(self, input_shape):
        return input_shape[0]

組み込み

以下の様にして、resnet構造を持つ任意のモデルで使用可能。

resblock.py

return layers.Add(
    [inputs, Shakedrop(num_of_unit=num_of_unit, num_of_layers=num_of_layers)(x)])

終わりに

弊社では現在人材募集中です！
機械学習をゼロから学べる研修付きのインターン募集など色々ありますので、ご興味ある方はお気軽にご連絡ください

Wantedly: https://www.wantedly.com/companies/company_5667111/projects
Homepage: https://www.algoage.net/
Twitter: https://twitter.com/algoage