基於TensorFlow構建圖片風格遷移模型

07-18

來自專欄機器不學習15 人贊了文章

作者：天雨粟
專欄：機器不學習
GitHub：NELSONZHAO
本篇文章中代碼地址：NELSONZHAO/zhihu

前言

這兩天湊巧看到了一些Image Style Transfer的內容，就看了下論文並且用TensorFlow跑了下模型。Style Transfer屬於遷移學習的一種，就是通過將一個圖片的風格（style）運用到另一個圖片上，有點類似現在各種P圖軟體中濾鏡的功能。

我們來看幾個模型跑的結果：

再來看下style transfer中圖像的生成過程：

上述過程是利用CNN通過一張content image和style image進行學習，最終得到具有style image中style風格的content image，例如上述三個例子顯示了合成圖片隨著模型的訓練而逐漸學習到特定的style，第一張是具有梵高《星月夜》風格的漫威圖片，第二張是具有畢加索《格爾尼卡》風格的漫威圖片，第三章是《吶喊》風格的漫威圖片。

本篇實戰代碼就將基於TensorFlow模型構建一個Style Transfer模型，實現圖片風格的遷移。

整篇文章的代碼參考來源於Stanford CS20課程的第二次作業內容
論文參考A Neural Algorithm of Artistic Style。

正文

正文部分主要分為四個部分：

Style Transfer原理簡介
VGG預訓練模型的載入
Style Transfer模型的構建
模型結果

本代碼運行環境：

Mac OS
Python 3.6
TensorFlow 1.6

如果有代碼報錯，請各位小夥伴先檢查是不是環境和版本問題。

本文代碼地址：NELSONZHAO/zhihu

一. Style Transfer原理簡介

由於本文主要側重實戰代碼，因此對於原理介紹也採用通俗易懂的方式進行闡述。Style Transfer即將一張圖片的風格應用到另一張圖片上，前者被稱為style image，後者被稱為content image。

通過content image和style image得到的圖片叫做合成圖片（synthetic image），它同時具有content image的內容信息和style image的風格信息。

因此我們要通過一個模型分別捕捉到content image的內容信息和style image的風格信息，在處理圖像問題中，CNN（卷積神經網路）被證明是一種提取圖像特徵的有效模型，因此我們採用CNN作為整個模型的信息提取。

整個模型的思路如下：

初始化合成圖（synthetic image）；
從Content Image中提取內容信息，最小化合成圖內容和Content Image內容的差異；
從Style Image中提取風格信息，最小化合成圖風格和Style Image風格的差異；
更新合成圖的信息，得到最終synthetic image。

初始化合成圖

由於我們最終是想將style image上的風格遷移到content image上，因此我們的合成圖最終在內容上與content image應該是一致的，例如裡面的物體、背景等信息。於是我們採用在content image上添加白雜訊的方式來初始化合成圖。

Content信息捕捉

Content信息來源於Content Image。對於Content Image來說，淺層卷積層的feature maps已經能夠重構出比較好的圖像了，而隨著卷積層的加深，從feature map重構出的圖像損失的像素點越嚴重，但是卻能夠保留圖像中的更加高級的內容信息，如下圖Content Restructions部分。可以看到隨著卷積層的加深，房屋圖片的像素點越來越模糊，但是仍然能看到整個圖片中的內容信息。

Style信息捕捉

Style信息來源於Style Image，在Style Image重構過程中，隨著卷積層的加深，feature maps重構出的圖像能夠更大程度上保留圖像的style信息，但是會丟失原始圖像中內容信息，例如重構出的圖像完全不具備原有圖像中物體的位置信息等。如上圖Style Reconstructions所示。

最小化loss

當模型能夠分別捕捉到content和style信息後，只需要對我們的合成圖進行學習，最小化content loss和style loss即可。

二. VGG預訓練模型載入

我們在此次模型中採用VGG-19作為pre-trained模型，我們只需要載入並獲取VGG的權重，在模型中進行feature maps的提取即可。

VGG-19模型下載地址。文件大小約為500M，裡面存儲了預訓練好的參數。

VGG-19的模型結構如下：

可以看到VGG-19中共有5段卷積，每段卷積中包含2-4個卷積層，並採用max-pooling作為池化層。

我們首先構造一個函數獲取指定層的權重信息：

在這個函數中，我們獲取指定層的權重數據。

接下來要構造我們的卷積層和池化層。在論文中，作者提到了average-pooling比max-pooling效果更好，因此，我們重新定義了一個average-pooling層。

在avgpool中，ksize和strides參數可以自行修改。

最後定義load函數，構造我們的模型graph，這裡不同於原始VGG的是，我們將原有的max-pooling替換成了average-pooling。代碼中convi_j代VGG中第i段的第j個卷積層，例如conv4_1代表第4段卷積中的第1個卷積層。

有了VGG中預訓練好的參數，我們可以載入進行在style transfer模型中使用，完整代碼在我的GitHub里load_vgg.py文件中。

三. Style Transfer模型的構建

模型的構建我們拆分為以下7個步驟：

載入VGG參數
計算content loss
計算style loss
計算總體loss
優化函數
構建模型圖
訓練模型

1. 載入VGG參數

調用我們前面寫好的load_vgg.py文件，構造VGG對象並且將參數載入進來。需要注意的是，VGG在圖片會將圖像進行mean-center操作，所以我們同樣也要對我們的content_img和style_img進行mean-center。

2. 計算content loss

在論文中，作者對於content loss的定義如下：

其中 $overrightarrow{p}$ 代表content image， $overrightarrow{x}$ 代表我們的合成圖片。 $P^l$ 代表content image在第 $l$ 層的feature representations，同樣的， $F^l$ 代表合成圖片在第 $l$ 層的feature representations。在不考慮batch size的情況下， $F,Pinmathbb{R}^{N_l imes M_l}$ ，其中 $N_l$ 代表第 $l$ 層的feature maps數量， $M_l$ 代表第 $l$ 層的feature maps的 $heights imes widths$ 。

但是，在實際學習過程中，上面這個content loss公式很難收斂，於是有很多人用 $frac{1}{4 imes P.size}$ 替代了 $frac{1}{2}$ ，其中 $P.size=width imes height imes channel$ 。

原論文中作者建議使用VGG中conv4_2作為content representation。

根據上面的公式，我們定義了如下函數來計算content loss：

3. 計算style loss

style loss的計算相比content loss要略複雜。在論文中，style loss公式定義如下：

我們首先看第一個公式，在上述公式中， $G^l$ 代表style image在第 $l$ 層feature maps的Gram Matrix（格雷姆矩陣）， $A^l$ 代表合成圖片在第 $l$ 層feature maps的Gram Matrix。 $N_l$ 代表第 $l$ 層feature maps的數量， $M_l$ 等於feature maps的 $width imes height$ 。Gram Matrix可以獲取到當前 $l$ 層中不同feature maps的feature correlations。

不同於content loss，style loss中對多個層的feature representation進行加權計算，得到總的style loss，如第二個公式所示，其中 $overrightarrow{a}$ 代表style image， $overrightarrow{x}$ 代表合成圖像， $w_l$ 代表第 $l$ 層的權重。

在論文中，作者使用了conv1_1，conv2_1，conv3_1，conv4_1，conv5_1。關於權重 $w_l$ ，可以自行調整，論文中使用了平均權重0.2，但我們知道higher layers提取到的style更好，所以在代碼中我使用了遞增的權重，即給更深的卷積層更大的權重。

根據上面的公式和原理，我們可以定義三個函數：

計算Gram Matrix
計算單個層的Style Loss
加權求和總的Style Loss

首先是Gram Matrix，將當前層的每個feature maps reshape為 $M imes N$ ，再進行Gram Matrix計算：

緊接著我們定義單層的style loss，即獲取Style Image和合成圖片的Gram Matrix，再按照上述公式計算Style loss：

最後，將多層的style loss進行加權求和得到總的style loss：

4. 計算總體loss

在有了合成圖片的content loss和style loss以後，我們就可以加權得到total loss。論文中公式如下：

即對合成圖片的content loss和style loss進行加權求和，一般 $frac{alpha}{eta}=0.001$ 或 $0.0001$ ，可以自己調整。模型總體loss代碼如下：

在這裡我們用到了一個trick，就是模型中的三張圖content image，style image和合成圖片共用了一個tensor：input_img，這是因為我們需要對這三張圖進行相同的操作，input_img的定義如下：

在使用過程中，我們通過tensorflow中的assign方法對變數input_img更改，從而實現共用同一個variable。

5. 優化函數

這裡優化函數採用AdamOptimizer：

6. 訓練模型

模型各個部分函數定義好以後，再構造一個build函數，將各個組件初始化：

進而構造train函數來訓練模型，在模型中還定義了一些summary和graph的記錄和存儲，具體完整代碼見GitHub。

這裡要注意的是，我們在使用VGG時候對圖像按照channel進行了mean-center，在這裡我們合成圖像要執行逆向操作，即在每個channel上加像素值的均值。

最終調用main函數執行函數：

我們指定content_img和style_img，並且指定圖像的像素尺寸（寬度和高度），在模型訓練之前，針對尺寸大小不同圖像會進行resize（具體代碼見GitHub），將content image，style image和合成圖片的像素尺寸限定在同樣的大小。

四. 模型結果

在代碼文件styles文件中，有一些style image的圖像，網上下載了一張漫威的圖片，對每種style訓練了300輪，得到如下結果：

還有風景圖：

星空圖：

總結

整個style transfer演算法的核心在於content和style信息的捕捉與loss function的設計，通過卷積得到的feature map去捕捉圖像的content和style，進而最小化合成圖像與原始圖像content和style的加權損失，最終學習得到具有style image風格的合成圖像。

專欄：機器不學習
個人簡介：天雨粟
GitHub：NELSONZHAO
轉載請聯繫作者獲得授權。