風格遷移

《從鍋爐工到AI專家(8)》中我們介紹了一個“圖片風格遷移”的例子。因為所引用的作品中使用了TensorFlow 1.x的程式碼，演算法也相對複雜，所以文中沒有仔細介紹風格遷移的原理。
今天在TensorFlow 2.0的幫助，和新演算法思想的優化下，實現同樣功能的程式碼量大幅減少，結構也越發清晰。所以今天就來講講這個話題。

“風格遷移”指的是將藝術作品的筆觸、技法等表現出來的視覺效果，應用在普通照片上，使得所生成的圖片，類似使用同樣筆觸、技法所繪製完成，但內容跟照片相同的“偽畫作”。
在神經網路機器學習的幫助下，生成圖片的觀賞性非常高，遠非早期傳統方法得到的圖片可比。
這裡重貼一遍前文中的例圖，讓我們有一個更直觀的感受。

首先是一張原程式作者的的自拍照：

接著不陌生，著名大作《星空》：

（請將以上兩圖儲存至工作目錄，不要修改檔名，我們稍晚的程式碼中會用到。）
兩張圖片經過程式處理後，會得到一幅新的圖片：

即使用《星空》風格模仿的手繪作品《黃粱一夢》:)

基本原理

風格遷移原理基於論文《A Neural Algorithm of Artistic Style》。
雖然論文中並沒有明說，但採用卷積神經網路做影象的風格遷移應當屬於一個實驗科學的成果而非單純的理論研究。
我們再引用一張前系列講解CNN時候的圖片：

一張圖片資料所形成的矩陣，在經過卷積網路的時候，影象中邊緣等視覺特徵會被放大、強化，從而形成一種特殊的輸出。通常我們只關心資料結果，並沒有把這些資料還原為圖片來觀察。而論文作者不僅這樣做了，恐怕還進行了大量的實驗。

這些神經網路中間結果圖片具有如此典型的特徵，可以脫離出主題內容而成為單純風格的描述。被敏銳的作者抓住深入研究也就不奇怪了。

最終研究成果確立了卷積神經網路進行圖片遷移的兩大基礎演算法：

• 在神經網路中，確定的抽取某些層代表內容的數字描述，以及另外一些層代表風格的數字描述。
  
• 多個層的輸出資料，通過公式的計算，擬合到同輸入影象相同的色域空間。這個公式即能用於代價函式中原始風格同目標風格之間的對比，也可以變形後通過組合多個風格層，生成新的目標圖片。

本系列文章都是盡力不出現數學公式，用程式碼講原理。
在《從鍋爐工到AI專家(8)》引用的程式碼中，除了構建神經網路、訓練，主要工作是在損失函式降低到滿意程度之後，使用網路中間層的輸出結果計算、組合成目標圖片。原文中對這部分的流程也做了簡介。

新的程式碼來自TensorFlow官方文件。除了程序升級為TensorFlow 2.0原生程式碼。在圖片的產生上也做了大幅創新：使用照片圖片訓練神經網路，每一階梯的訓練結果，不應用回神經網路（網路的權重引數一直固定鎖死的），而把訓練結果應用到圖片本身。在下一次的訓練迴圈中，使用新的圖片再次計算損失值。這樣，當損失值最小的時候，訓練圖片本身就已經是符合我們要求的生成圖片。當然本質上，跟前一種方法一樣的。但感覺上，結構清晰了很多。這個過程對比起來，大量節省了圖片生成的計算。當然，主要原因還是TensorFlow 2.0內建的tf.linalg.einsum方法強大好用。

在特徵層的定義上，照片內容的描述使用vgg-19網路的第5部分的第2層卷積輸出結果。藝術圖片風格特徵的描述使用了5個層，分別是vgg-19網路的第1至第5部分第1個網路層的輸出結果。在程式中，可以這樣描述：

# 定義最能代表內容特徵的網路層
content_layers = ['block5_conv2'] 

# 定義最能代表風格特徵的網路層
style_layers = ['block1_conv1',
                'block2_conv1',
                'block3_conv1', 
                'block4_conv1', 
                'block5_conv1']

網路層的名稱來自於vgg-19網路定義完成後，各層的名稱。可以使用如下程式碼得到所有層的名稱：

    ...
# 建立無需分類結果的vgg網路
vgg = tf.keras.applications.VGG19(include_top=False, weights='imagenet')

# 顯示vgg中所有層的名稱
print()
for layer in vgg.layers:
    print(layer.name)
    ...

通常的模型訓練，都是使用代價函式比較網路輸出結果，和目標標註值的差異，使得差異逐漸縮小。
本例的訓練目標比較複雜，可以描述為兩條：

• 生成圖片的風格層輸出，同藝術圖片的風格層輸出差異最小
• 生成圖片的內容層輸出，同原始照片的內容層輸出差異最小化

雖然這個代價函式略微複雜，不過比VAE的代價函式還是簡單多了:)

原始碼

程式中的註釋非常詳細。跟以前的程式有一點區別，就是直接使用TensorFlow內建方法讀取了圖片檔案，然後呼叫jpg解碼還原為矩陣。
不過TensorFlow內建的將影象0-255整數值轉換為浮點數的過程，會自動將數值變為0-1的浮點小數。
這個過程其實對我們多此一舉，因為我們後續的很多計算都需要轉換回0-255。

#!/usr/bin/env python3

from __future__ import absolute_import, division, print_function, unicode_literals

import tensorflow as tf
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
import time
import functools
import time
from PIL import Image

# 設定繪圖視窗引數，用於圖片顯示
mpl.rcParams['figure.figsize'] = (13, 10)
mpl.rcParams['axes.grid'] = False

# 獲取下載後本地圖片的路徑，content_path是真實照片，style_path是藝術品風格圖片
content_path = "1-content.jpg"
style_path = "1-style.jpg"

# 讀取一張圖片，並做預處理
def load_img(path_to_img):
    max_dim = 512
    # 讀取二進位制檔案
    img = tf.io.read_file(path_to_img)
    # 做JPEG解碼，這時候得到寬x高x色深矩陣，數字0-255
    img = tf.image.decode_jpeg(img)
    # 型別從int轉換到32位浮點，數值範圍0-1
    img = tf.image.convert_image_dtype(img, tf.float32)
    # 減掉最後色深一維，獲取到的相當於圖片尺寸（整數），轉為浮點
    shape = tf.cast(tf.shape(img)[:-1], tf.float32)
    # 獲取圖片長端
    long = max(shape)
    # 以長端為比例縮放，讓圖片成為512x???
    scale = max_dim/long
    new_shape = tf.cast(shape*scale, tf.int32)
    # 實際縮放圖片
    img = tf.image.resize(img, new_shape)
    # 再擴充套件一維，成為圖片數字中的一張圖片（1，長，寬，色深）
    img = img[tf.newaxis, :]
    return img

# 讀入兩張圖片
content_image = load_img(content_path)
style_image = load_img(style_path)

############################################################
# 定義最能代表內容特徵的網路層
content_layers = ['block5_conv2'] 

# 定義最能代表風格特徵的網路層
style_layers = ['block1_conv1',
                'block2_conv1',
                'block3_conv1', 
                'block4_conv1', 
                'block5_conv1']
# 神經網路層的數量
num_content_layers = len(content_layers)
num_style_layers = len(style_layers)

# 定義一個工具函式，幫助建立得到特定中間層輸出結果的新模型
def vgg_layers(layer_names):
    """ Creates a vgg model that returns a list of intermediate output values."""
    # 定義使用ImageNet資料訓練的vgg19網路
    vgg = tf.keras.applications.VGG19(include_top=False, weights='imagenet')
    # 已經經過了訓練，所以鎖定各項引數避免再次訓練
    vgg.trainable = False
    # 獲取所需層的輸出結果
    outputs = [vgg.get_layer(name).output for name in layer_names]
    # 最終返回結果是一個模型，輸入是圖片，輸出為所需的中間層輸出
    model = tf.keras.Model([vgg.input], outputs)
    return model

# 定義函式計算風格矩陣，這實際是由抽取出來的5個網路層的輸出計算得來的
def gram_matrix(input_tensor):
    result = tf.linalg.einsum('bijc,bijd->bcd', input_tensor, input_tensor)
    input_shape = tf.shape(input_tensor)
    num_locations = tf.cast(input_shape[1]*input_shape[2], tf.float32)
    return result/(num_locations)

# 自定義keras模型
class StyleContentModel(tf.keras.models.Model):
    def __init__(self, style_layers, content_layers):
        super(StyleContentModel, self).__init__()
        # 自己的vgg模型，包含上面所列的風格抽取層和內容抽取層
        self.vgg = vgg_layers(style_layers + content_layers)
        self.style_layers = style_layers
        self.content_layers = content_layers
        self.num_style_layers = len(style_layers)
        # vgg各層引數鎖定不再引數訓練
        self.vgg.trainable = False

    def call(self, input):
        # 輸入的圖片是0-1範圍浮點，轉換到0-255以符合vgg要求
        input = input*255.0
        # 對輸入圖片資料做預處理
        preprocessed_input = tf.keras.applications.vgg19.preprocess_input(input)
        # 獲取風格層和內容層輸出
        outputs = self.vgg(preprocessed_input)
        # 輸出實際是一個數組，拆分為風格輸出和內容輸出
        style_outputs, content_outputs = (
                outputs[:self.num_style_layers],
                outputs[self.num_style_layers:])
        # 計算風格矩陣
        style_outputs = [gram_matrix(style_output)
                         for style_output in style_outputs]

        # 轉換為字典
        content_dict = {content_name: value
                        for content_name, value
                        in zip(self.content_layers, content_outputs)}
        # 轉換為字典
        style_dict = {style_name: value
                      for style_name, value
                      in zip(self.style_layers, style_outputs)}
        # 返回內容和風格結果
        return {'content': content_dict, 'style': style_dict}

# 使用自定義模型建立一個抽取器
extractor = StyleContentModel(style_layers, content_layers)

# 設定風格特徵的目標，即最終生成的圖片，希望風格上儘量接近風格圖片
style_targets = extractor(style_image)['style']
# 設定內容特徵的目標，即最終生成的圖片，希望內容上儘量接近內容圖片
content_targets = extractor(content_image)['content']

# 內容圖片轉換為張量
image = tf.Variable(content_image)

# 擷取0-1的浮點數，超範圍部分被擷取
def clip_0_1(image):
    return tf.clip_by_value(image, clip_value_min=0.0, clip_value_max=1.0)

# 優化器
opt = tf.optimizers.Adam(learning_rate=0.02, beta_1=0.99, epsilon=1e-1)
# 預定義風格和內容在最終結果中的權重值，用於在損失函式中計算總損失值
style_weight = 1e-2
content_weight = 1e4

# 損失函式
def style_content_loss(outputs):
    style_outputs = outputs['style']
    content_outputs = outputs['content']
    # 風格損失值，就是計算方差
    style_loss = tf.add_n([tf.reduce_mean((style_outputs[name]-style_targets[name])**2) 
                           for name in style_outputs.keys()])
    # 權重值平均到每層，計算總體風格損失值
    style_loss *= style_weight/num_style_layers

    # 內容損失值，也是計算方差
    content_loss = tf.add_n([tf.reduce_mean((content_outputs[name]-content_targets[name])**2) 
                             for name in content_outputs.keys()])
    content_loss *= content_weight/num_content_layers
    # 總損失值
    loss = style_loss+content_loss
    return loss
################################################################

# 一次訓練
@tf.function()
def train_step(image):
    with tf.GradientTape() as tape:
        # 抽取風格層、內容層輸出
        outputs = extractor(image)
        # 計算損失值
        loss = style_content_loss(outputs)

    # 梯度下降
    grad = tape.gradient(loss, image)
    # 應用計算後的新引數，注意這個新值不是應用到網路
    # 作為訓練完成的vgg網路，其引數前面已經設定不可更改
    # 這個引數實際將應用於原圖
    # 以求取，新圖片經過網路後，損失值最小
    opt.apply_gradients([(grad, image)])
    # 更新圖片，用新圖片進行下次訓練迭代
    image.assign(clip_0_1(image))

start = time.time()
epochs = 10
steps_per_epoch = 100

step = 0
for n in range(epochs):
    for m in range(steps_per_epoch):
        step += 1
        train_step(image)
        print(".", end='')
    print("")
    # 每100次迭代顯示一次圖片
    # imshow(image.read_value())
    # plt.title("Train step: {}".format(step))
    # plt.show()

end = time.time()
print("Total time: {:.1f}".format(end-start))

########################################
#儲存結果圖片
file_name = 'newart1.png'
mpl.image.imsave(file_name, image[0])

程式的輸出結果如下圖：

看起來基本達到了設計要求，不過再仔細觀察，似乎效果雖然都有了，但畫面看上去有一點不乾淨，有很多小的噪點甚至有了干涉紋。
這是因為，在照片原圖和藝術作品原圖中，肯定天然就存在有噪點以及圖片中本身應當有的小而頻繁的花紋。這些內容在通過卷積加強後，兩幅照片再疊加，這些噪聲就被強化了，從而在生成的圖片中體現的非常明顯。
這個問題如果在傳統演算法中可以使用高通濾波。在卷積神經網路中則更容易，是統計總體變分損失值(Total Variation Loss)，在代價函式中，讓這個損失值降到最小，就抑制了這種噪點的產生。也相當於神經網路具有了降噪的效果。
變分損失是計算圖片中，在X方向及Y方向，相鄰畫素的差值。如果畫素差別不大，那差肯定很小甚至趨近於0。如果差別大，當然差值就大。
請使用下面的程式碼，替換上面程式中訓練的部分：

###################################################
# 計算x方向及y方向相鄰畫素差值，如果有高頻花紋，這個值肯定會高，
# 因為相鄰點相同差值接近0，區別越大，差值當然越大
def high_pass_x_y(image):
    x_var = image[:, :, 1:, :] - image[:, :, :-1, :]
    y_var = image[:, 1:, :, :] - image[:, :-1, :, :]

    return x_var, y_var

# 計算總體變分損失
def total_variation_loss(image):
    x_deltas, y_deltas = high_pass_x_y(image)
    return tf.reduce_mean(x_deltas**2)+tf.reduce_mean(y_deltas**2)


# 總體變分損失值在損失值中所佔權重
total_variation_weight = 1e8

# 一次訓練
@tf.function()
def train_step(image):
    with tf.GradientTape() as tape:
        # 抽取風格層、內容層輸出
        outputs = extractor(image)
        # 計算損失值
        loss = style_content_loss(outputs)
        loss += total_variation_weight*total_variation_loss(image)

    # 梯度下降
    grad = tape.gradient(loss, image)
    # 應用計算後的新引數，注意這個新值不是應用到網路
    # 作為訓練完成的vgg網路，其引數前面已經設定不可更改
    # 這個引數實際將應用於原圖
    # 以求取，新圖片經過網路後，損失值最小
    opt.apply_gradients([(grad, image)])
    # 更新圖片，用新圖片進行下次訓練迭代
    image.assign(clip_0_1(image))

# 內容圖片作為逐步迭代生成的新圖片，一開始當然是原圖，這裡是轉換為張量
image = tf.Variable(content_image)

start = time.time()

# 迭代10次，每次100步訓練
epochs = 10
steps = 100

step = 0
for n in range(epochs):
    for m in range(steps):
        step += 1
        train_step(image)
        print(".", end='')
    print("")
end = time.time()
print("Total time: {:.1f}".format(end-start))

#儲存結果圖片
file_name = 'newart1.png'
mpl.image.imsave(file_name, image[0])

再次執行，所得到的輸出圖片如下：

效果不錯吧？可以換上自己的照片還有自己心儀的藝術作品來試試。
程式中限制了圖片寬、高最大值是512，如果裝置效能比較好，或者有更大尺寸的需求，可以修改程式中的常量。

（待續...