Kerasで自作のU-netを用いて画像のスタイル変換をしようとしましたが、Loss関数をMSEで学習しても、スタイルをうまく学習できませんでした。
そこで、スタイル変換の論文でよく使われているPerceptual Lossを使ってみることにしました。

Perceptual Loss

Perceptual Lossとは、画像を学習済みネットワークに通して得られる特徴マップ同士でlossを計算します。画像同士でのMSEでは、ピクセル単位でlossが発生してしまい、スタイルがうまく学習できなかったり、ぼやけた出力になってしまったりします。それに対し、特徴マップ同士でのlossでは、ピクセルのずれでlossが発生しにくいので、スタイルの学習やシャープな出力が期待されます。
また、ネットワークの最後の層の出力だけでなく、途中の層の特徴マップもlossの計算に使うことで、様々な解像度のスタイルで学習を行うことができます。
学習済みのモデルとしてVGG-16などがよく使われているようですが、このモデルは特徴抽出器の後に全結合層のクラス分類器が接続されています。Perceptual Lossの計算に必要なのは特徴マップのみなので、全結合層は使わず特徴抽出器のみを使います。

実装

U-NetなどのAuto Encoderのネットワークを用いて、画像のスタイル変換を行う場合の実装です。スタイル変換のネットワークには、自作またはimportしたモデルを使うことを前提としています。

学習済みVGG-16

今回はKerasのApplicationsで様々な学習済みモデルが提供されているので、その中のVGG-16を使います。今回は全結合層は必要ないので、モデルを作る際に引数に

include_top=False

を指定することで、特徴抽出器のみになります。また、この場合、モデルに入力するデータ形状を指定することができます。

from keras.models import Model
from keras.applications.vgg16 import VGG16

# Kerasの学習済みVGG-16
vgg_model = VGG16(weights='imagenet', # ImageNetで学習したモデル
                  include_top=False, # 特徴抽出器のみ
                  input_shape=(512, 512, 3)) # 入力形状
vgg_model.trainable = False # 学習時にVGGのパラメータ更新をしない

loss_model = Model(vgg_model.input, vgg_model.output) # Loss計算のためのモデルを定義
loss_model.compile(optimizer='adam', loss='mse', metrics=['accuracy'])

スタイル変換を含めたネットワーク

上記で実装したLoss計算モデルに、スタイル変換モデルを追加します。

transfer_model = unet_model(shape=(512, 512, 3)) # 512x512,3chの画像を変換するモデル
loss_model_outputs = loss_model(transfer_model.output) # transfer_modelの出力をloss_modelに入力
full_model = Model(transfer_model.input, loss_model_outputs) # 全体のネットワーク
full_model.compile(optimizer='adam', loss='mse', metrics=['accuracy']) # 特徴マップをMSEでLoss計算

unet_model(shape=(512, 512, 3))

は、自作のモデル定義関数なので、適切なものに変更してください。

学習&テスト

上記で実装したモデルの出力は、VGGの特徴マップになっていますが、目的の出力はtransfer_modelで変換した画像になります。そこで、モデルに入力する教師データと、テスト時の出力を工夫する必要があります。
以下のデータがあることを前提としています。

x_train : 512x512x3の画像が複数枚あるnumpy array
y_train : x_trainに対応する教師画像
x_test : 512x512x3の画像が複数枚あるnumpy array
y_test : x_testに対応する教師画像

学習

実装したfull_modelの入力は画像、出力はloss_modelの特徴量であるので、full_modelに与える教師データもloss_model特徴量に変換します。

y_train_feature = loss_model.predict(y_train) # 教師画像を特徴量に変換
full_model.fit(x_train, y_train_feature, epoch=10)

テスト

テスト時の出力は、full_modelの中のtransfer_modelの出力とします。

full_model.summary()

で、モデルの層を調べて、対応する層の出力をpred_modelの出力とします。

pred_model = Model(full_model.input, full_model.layers[50].output) # 数字はtransfer_modelの層数に応じて変更
predict = pred_model.predict(x_test)

補足

複数の層でLossを計算

VGGの最後の出力でLossを計算しましたが、複数の層の出力でLossを計算するには、loss_modelのところで以下のコードを追加・変更します。

##### 追加
select_layers = [2, 9, 18] # Lossを計算する層の選択
vgg_outputs = [vgg_model.layers[i].output for i in select_layers] # 選択した層の出力
####

loss_model = Model(vgg_model.input, vgg_outputs) # 変更

画像のMSEもLossに追加する

特徴マップでのLossだけでなく、画像のピクセル単位のMSEもLossに追加したいです。

vgg_outputs.append(vgg_model.input)

で、vgg_modelの入力（transfer_modelの出力画像）をloss_modelの出力に加えることで、画像のLossも追加しようとしましたが、「モデルの入力を出力にできない」エラーが出てしまい、できませんでした。
そこで、画像の値は正であるので、入力にReLUを通して、その出力をloss_modelの出力に加えることで対処しました。（もっといい対処法があると思います）

from keras.layers import Input, Activation

#### 追加
input_batch = Input(shape=(512, 512, 3))
relu_output = Activation('relu')(input_batch)
####

#### 変更
vgg_model = VGG16(weights='imagenet', # ImageNetで学習したモデル
                  include_top=False, # 特徴抽出器のみ
                  input_tensor=input_batch, # 追加
                  input_shape=(512, 512, 3)) # 入力形状
####

vgg_outputs.append(relu_output) # 追加
loss_model = Model(input_batch, vgg_outputs) # 変更

参考にしたサイト

Keras Applications

https://keras.io/ja/applications/

Implement perceptual loss with pretrained VGG using keras

https://stackoverflow.com/questions/47675094/implement-perceptual-loss-with-pretrained-vgg-using-keras