スタイル変換のStyle Lossとは何をやっているか

スタイル変換やImage to Imageの損失関数で使われる・Style Lossの実装を詳しく見ていきます。Style Lossの計算で用いているグラム行列の計算方法をTensorFlowで考えます。

Style Lossのやっていること

2つの画像の、VGG16や19（どっちを使うか、どのレイヤーを使うかは論文によって異なる）の中間層の値を取り出し、ぞれぞれのグラム行列を取りL1ロスを計算する。

比較する画像は論文によって異なるが、例えばP-Convの場合は、復元画像とGround Truthの画像を比較する。

実装から見る

例えばStyle Lossを使っているP-Convの場合。コードはこちらより。

https://github.com/MathiasGruber/PConv-Keras/blob/master/libs/pconv_model.py

    def loss_style(self, output, vgg_gt):
        """Style loss based on output/computation, used for both eq. 4 & 5 in paper"""
        loss = 0
        for o, g in zip(output, vgg_gt):
            loss += self.l1(self.gram_matrix(o), self.gram_matrix(g))
        return loss

グラム行列のL1ロスを取っています。L1ロスはおなじみの定義ですね。

    @staticmethod
    def l1(y_true, y_pred):
        """Calculate the L1 loss used in all loss calculations"""
        if K.ndim(y_true) == 4:
            return K.mean(K.abs(y_pred - y_true), axis=[1,2,3])
        elif K.ndim(y_true) == 3:
            return K.mean(K.abs(y_pred - y_true), axis=[1,2])
        else:
            raise NotImplementedError("Calculating L1 loss on 1D tensors? should not occur for this network")

で、問題はグラム行列のほうです。

    @staticmethod
    def gram_matrix(x, norm_by_channels=False):
        """Calculate gram matrix used in style loss"""

        # Assertions on input
        assert K.ndim(x) == 4, 'Input tensor should be a 4d (B, H, W, C) tensor'
        assert K.image_data_format() == 'channels_last', "Please use channels-last format"        

        # Permute channels and get resulting shape
        x = K.permute_dimensions(x, (0, 3, 1, 2))
        shape = K.shape(x)
        B, C, H, W = shape[0], shape[1], shape[2], shape[3]

        # Reshape x and do batch dot product
        features = K.reshape(x, K.stack([B, C, H*W]))
        gram = K.batch_dot(features, features, axes=2)

        # Normalize with channels, height and width
        gram = gram /  K.cast(C * H * W, x.dtype)

        return gram

入力が$(B, H, W, C)$という4階テンソルだとしましょう。Bはバッチサイズ、Hは縦解像度、Wは横解像度、Cはチャンネル数を表します。

コードを順におっていくと、まずは次元を入れ替えます。$(B, C, H, W)$としています。次に、$(B, C, HW)$という3階に変形します。そのあとのbatch_dotというのがKeras独特の関数なのですが、1階以降に対して、

$$(C, HW)\dot (HW, C)=(C, C)$$

という計算をするのがbatch_dotの関数です。

batch_dotをTensorFlowの関数で置き換えて考える

batch_dotをもう少しわかりやすい関数で見てみます。以下の処理とbatch_dotは同じ結果を出します。

import tensorflow as tf
import tensorflow.keras.backend as K

def batch_dot():
    x = tf.reshape(tf.range(24), (4, 2, 3))
    # Kerasのbatch_dot
    a = K.batch_dot(x, x, axes=2)
    print("Keras batch dot")
    print(a)
    # tfの関数で再現
    y = tf.transpose(x, (0, 2, 1))
    b = tf.matmul(x, y)
    print("Implemented by tensorflow function")
    print(b)
    print("DIff")
    print(b - a)

if __name__ == "__main__":
    batch_dot()

結果は次の通り。

Keras batch dot
tf.Tensor(
[[[   5   14]
  [  14   50]]

 [[ 149  212]
  [ 212  302]]

 [[ 509  626]
  [ 626  770]]

 [[1085 1256]
  [1256 1454]]], shape=(4, 2, 2), dtype=int32)
Implemented by tensorflow function
tf.Tensor(
[[[   5   14]
  [  14   50]]

 [[ 149  212]
  [ 212  302]]

 [[ 509  626]
  [ 626  770]]

 [[1085 1256]
  [1256 1454]]], shape=(4, 2, 2), dtype=int32)
DIff
tf.Tensor(
[[[0 0]
  [0 0]]

 [[0 0]
  [0 0]]

 [[0 0]
  [0 0]]

 [[0 0]
  [0 0]]], shape=(4, 2, 2), dtype=int32)

この通り、計算結果が同一になります。つまり、3階に変形して、1階以降を転置したものとの積を取ったのがbatch_dotなのです。

一般にグラム行列といえば、2階の行列に対して$WW^T$を計算するのがグラム行列（左右は逆転してもOK）なので、TensorFlowで実装したコードは3階に拡張したグラム行列の定義そのものになります。

グラム行列の直感的な解釈

そもそも「なんでスタイルの損失を決めるのにグラム行列を取るの？」ということなのですが、分散共分散行列・相関行列はグラム行列の特別な場合です。詳しくはこちらの記事で書いたのでを参照してみてください。

つまり、「VGGでの中間層での空間での、ピクセル間の相関のようなものを似せる」というのが、Style Lossのやっていることです。もしこれをグラム行列のL1ではなくピクセル間のL1を取るとPerceptual Lossというものになりますが、Perceptual LossとStyle Lossの違いとは、Style Lossのほうが相関を見ているので、より広範囲、より位置に依存しない特徴を見ているということになります。