機械学習奮闘記

CV系の論文実装とか

StyleGAN: A Style-Based Generator Architecture for Generative Adversarial Networks

論文

A Style-Based Generator Architecture for Generative Adversarial Networks(Karras et al., 2019)
いつもインパクト大な高解像度の画像生成手法を発表してくるNVIDIAの論文である.

概要

スタイル変換の文献を参考に,敵対的生成ネットワークにおけるジェネレータの新しいアーキテクチャを提案する. このアーキテクチャは,高レベルな属性(例:人の顔で訓練されたときのポーズ,アイデンティティなど)を教師なしで分離し, 生成画像に確率的な変動(例、そばかす,髪など)を含めることを可能にする. これにより,スケールに応じた直感的な画像生成の制御を可能にした. この新しいジェネレータは,従来の分布の品質の評価指標におけるSOTAを向上させ, 明らかに優れた補間特性をもたらし,また変動の潜在的要因をより良く解きほぐす. 補間の品質と解きほぐれ具合を定量化するために,我々は任意のジェネレータに適用可能な2つの新しい評価指標を提案する. 最後に,我々は非常に多様で高品質な,新しい人の顔のデータセットを紹介する.

今のところ,この潜在空間におけるもつれというのがあんまりイメージできないが,この辺りがキーになるのかな.

実装

完全なソースコードはここにある.フレームワークはTensorFlow. NVIDIA本家の実装が公開されているので,所々参考にしたが,読み解くのが結構辛かった. github.com

Style-Based Generator

f:id:sakuma_dayo:20190306195503p:plain

Mapping Network

StyleGANにおいては潜在ベクトル z \in Zはネットワークの入力ではなく,各レイヤーにおいてスタイル制御のために用いられる. 潜在ベクトルはMapping Networkによってスタイル制御のための潜在空間へマッピングされる.( f : Z \rightarrow W) Mapping Networkは論文では8層のMLPとして実装されている.

def mapping_network(latents, labels, reuse=tf.AUTO_REUSE):
    with tf.variable_scope("mapping_network", reuse=reuse):
        labels = embedding(
            inputs=labels,
            units=latents.shape[1],
            variance_scale=1,
            scale_weight=True
        )
        latents = tf.concat([latents, labels], axis=1)
        latents = pixel_norm(latents)
        for i in range(self.mapping_layers):
            with tf.variable_scope("dense_block_{}".format(i)):
                with tf.variable_scope("dense".format(i)):
                    latents = dense(
                        inputs=latents,
                        units=latents.shape[1],
                        use_bias=True,
                        variance_scale=2,
                        scale_weight=True
                    )
                latents = tf.nn.leaky_relu(latents)
        return latents

Adaptive Instance Normalization

Mapping Networkによって得られた w \in Wは各畳み込み層の後にアフィン変換されたのち, Adaptive Instance Normalization(AdaIN)(Hung et al., 2017)のパラメータとして用いられる. AdaINは以下の式で与えられる.

  \displaystyle AdaIN(x_{i}, y_{i})=y_{i, s} \frac{x_{i} - \mu(x_{i})}{\sigma(x_{i})} + y_{i, b}

ここで x_{i}は特徴マップ, y_{i}はアフィン変換された潜在ベクトルである. これはInstance Normalizationによって特徴マップ毎に正規化した後,スタイル変換のためのスケールとバイアスをMLPでモデル化している. これを各畳み込み層の後に行うことで,各スケール毎にスタイルを変化させることができる. ここでいうスケールは特徴の意味的なレベルとも捉えられる.

AdaINによるスタイル制御はスケールに対してローカルである.つまりあるスケールにおけるスタイルの変更は他のスケールのスタイルに影響を及ぼさない. これはAdaIN操作は,その後に続く畳み込みのために,特徴マップ間の相対的な重要性を変更するが,正規化処理により,元の特徴マップの統計には依存しないからである.

def adaptive_instance_norm(inputs, latents, use_bias=True, center=True, scale=True,
                           variance_scale=2, scale_weight=True, epsilon=1e-8):
    ''' Adaptive Instance Normalization
        [Arbitrary Style Transfer in Real-time with Adaptive Instance Normalization]
        (https://arxiv.org/pdf/1703.06868.pdf)
    '''
    # standard instance normalization
    inputs -= tf.reduce_mean(inputs, axis=[2, 3], keepdims=True)
    inputs *= tf.rsqrt(tf.reduce_mean(tf.square(inputs), axis=[2, 3], keepdims=True) + epsilon)

    if scale:
        with tf.variable_scope("scale"):
            gamma = dense(
                inputs=latents,
                units=inputs.shape[1],
                use_bias=use_bias,
                variance_scale=variance_scale,
                scale_weight=scale_weight
            )
            gamma = tf.reshape(
                tensor=gamma,
                shape=[-1, gamma.shape[1], 1, 1]
            )
        inputs *= gamma

    if center:
        with tf.variable_scope("center"):
            beta = dense(
                inputs=latents,
                units=inputs.shape[1],
                use_bias=use_bias,
                variance_scale=variance_scale,
                scale_weight=scale_weight
            )
            beta = tf.reshape(
                tensor=beta,
                shape=[-1, beta.shape[1], 1, 1]
            )
        inputs += beta

    return inputs

このAdaIN,cGANs with Projection Discriminator(Miyato et al. 2018),Spectral Normalization(Miyato et al. 2018)以降でよく用いられている Conditional Batch Normalization(de Vries et al., 2017)ともかなり近い. Conditional Batch Normalizationでは正規化後のスケールとバイアスをクラスラベルのembeddingでモデル化する.

Noise Inputs

StyleGANでは生成画像に確率的な要因を含めるために,各レイヤーにおいてノイズを供給する. これは1チャンネルのノイズマップを用意し,各特徴マップに対して個別にスケーリングした後,足し合わせる. このチャンネル数分のスケーリング係数はパラメータとして最適化される.

def apply_noise(inputs):
    noise = tf.random_normal([tf.shape(inputs)[0], 1, *inputs.shape[2:]])
    weight = tf.get_variable(
        name="weight",
        shape=[inputs.shape[1]],
        initializer=tf.initializers.zeros()
    )
    weight = tf.reshape(weight, [1, -1, 1, 1])
    inputs += noise * weight
    return inputs

これにより髪の毛,しわなどの確率的とみなせる多くの特徴をモデル化できる. また各レイヤー毎にノイズを供給するので,特徴のレベルに応じた確率的変動を実現できる. ただあくまでこれは低レベルな特徴を担い,高レベルな特徴はそのまま残る. (ように学習されることが期待できるって感じだと思う.高レベルな特徴までランダム化すると,それはもう人の顔と認識できなそう.)

AdaINによるスタイル制御は特徴マップ単位でスケール,バイアスを適用するので,その効果は画像全体にわたる効果を制御できる. 対してノイズはピクセル単位で供給されるので確率的な変動を制御できる.

f:id:sakuma_dayo:20190306195539p:plain

Constant input

StyleGANでは潜在ベクトルはもはやネットワークの入力である必要はなく,代わりにあらかじめ最小解像度のマップ用意しておき,常にこれをネットワークの入力とする. このマップはパラメータとして最適化される.

Style Mixing

まずMapping Networkに2つの潜在ベクトル z_{1}, z_{2} \in Zを入力し, w_{1}, w_{2} \in Wを得る. generatorのあるレイヤーまではAdaINのパラメータとして w_{1}を用いて,そこから先のレイヤーではAdaINのパラメータとして w_{2}を用いる. これはスタイルのスイッチングであり, w_{1}は高レベルのスタイルを, w_{2}は低レベルのスタイルを担っていると考えられる. スイッチするレイヤーは毎回ランダムに選ばれる. この正則化によって隣接するスタイルが相関していると,ネットワークが仮定するのを防ぐとあり, なんとなくそんな気がするも,ここがまだ少し飲み込めていない.

下図のように潜在ベクトルをコピーすることはスタイルをコピーすることであり,様々なレベルでスタイルを操作することが可能である.

f:id:sakuma_dayo:20190306195558p:plain

Low-Pass Filtering

本論文ではさらっと述べられている程度のことだが,少し気になってのでメモしておく. 本家の実装を見るとblur2dとかいう見慣れない関数があったので,論文を見てみるとnearest-neighbor samplingをローパスフィルタを用いたbilinear samplingに置き換えたと書いてある. これはICLR2019に投稿されているMaking Convolutional Networks Shift-Invariant Againで提案されている, ダウンサンプリング前にローパスフィルタをかけておくことで,シフト不変性を保つアプローチらしい. ダウンサンプリングではナイキスト周波数以上の高周波成分がエイリアシングとなって現れるので,ローパスフィルタによって高周波成分を除去することで,シフト不変性を保っている.

f:id:sakuma_dayo:20190306195800p:plain

音とかではよくあるアプローチだと思うが,画像では聞いたことなかったので面白かった. ただそこまで本質的というわけでもなさそうなのでとりあえず実装は後回しにする.

Zero-Gradient Penalty

StyleGANではWhich Training Methods for GANs do actually Converge?で提案されている Zero-Centered Gradient Penaltyを正則化項として用いている. この論文によるとWGAN(Arjovsky et al., 2017),WGAN-GP(Gulrajani et al., 2017)は必ずしも収束しないらしく, Zero-Centered Gradient Penaltyは常に収束するらしい. ちなみにWGAN-GPはOne-Centered Gradient Penaltyである. 論文では以下で示される2種類のZero-Centered Gradient Penaltyが提案されている.

  \displaystyle R_{1}(\psi) := \frac{\gamma}{2} E_{p_{\mathcal D}(x)}[\|\nabla D_{\psi}(x)\|^{2} ]

  \displaystyle R_{2}(\theta, \psi) := \frac{\gamma}{2} E_{p_{\theta}(x)}[\|\nabla D_{\psi}(x)\|^{2} ]

これらはそれぞれデータ分布,generator分布に対するgradient penaltyである. StyleGANではlossはNon-Saturating Loss(Goodfellow et al., 2014)とし, R_{1}-regularizerのみを用いている.

とりあえず今日はここまでにする.一番重要そうな潜在空間におけるもつれについて何も書いてないが, まだあんまり理解できてないので,完全に飲み込めたらまた書こうと思う. 早く実験結果を出したいところだが,この前実装したGANSynthに計算資源を食われているので, 今までめんどくさそうで避けていたGoogle Colaboratoryを使ってみた.このノートブックとかいう概念があんまり好きになれない...