從一次 CycleGAN 實現聊聊 TF

雷鋒網 AI科技評論按，本文作者Coldwings，該文首發于知乎專欄為愛寫程序，雷鋒網 AI科技評論獲其授權轉載。以下為原文內容，有刪減。

CycleGAN是個很有趣的想法（Unpaired Image-to-Image Translationusing Cycle-Consistent Adversarial Networks [https://arxiv.org/pdf/1703.10593.pdf]），看完這篇論文之后，隱隱地覺得，這后面有更多的內容可以挖，我盡我所能做出了各種嘗試，努力發掘更多的可能性。

實現過程中可以說還是略微糾結的，最初是用Keras快速實踐了一下，然而其實并不『快速』，后來反倒是用TensorFlow重寫以及嘗試各種意外想法時才感覺，當需要處理一些比較復雜的網絡結構、訓練流程甚至op時，TF提供的可以細化到每個操作的體驗實際上要比各種上層API都來得更好，而結合TensorBoard，可視化的訓練將取得更好的效果。當然，我對Torch無感，或許用Torch能有更好的體驗，但我不擅長這個；Chainer（A flexible framework for neural networks）講道理寫出來的代碼會更好看，但是似乎身邊用的人并不多，姑且放過。

這篇文章倒不是來介紹什么是CycleGAN的，若是不甚了解，我妻子將會將她的發表整理一下再發布出來（CycleGAN（以及DiscoGAN和DualGAN）簡介 - 知乎專欄）。這一陣的嘗試中，我自己也對GAN，對Generator中的圖像甚至其它東西的生成，以及單純從寫代碼角度來看，怎么管理TF里的變量，怎么把代碼寫得好看，怎么更好地利用TensorBoard都有了更多地理解，算是不小的提高吧……

所以這里也就大概提一提一些實現中需要注意的小技巧吧。（雖然我覺得大概大多數真正拿著TF搞DL研究的人都不需要研究這篇文章）

CycleGAN比較麻煩的地方

其實CycleGAN麻煩的地方不少，這是一個挺復合的模型：兩個Generator，兩個Discriminator，這已經是四個比較簡單的網絡了（是的，考慮到所有可能性，Generator和Discriminator完全可以各自都有兩種不同的結構）；一組Generator+Discriminator復合成一個GAN，又一層復合模型，并且GAN的訓練還得控制，由于G和D的損失相反，訓練G時需要控制D的變量讓其不可訓練；我們還要讓Cycle loss作為模型loss的一部分，這個更高一層的復合模型由兩個GAN組成……

良好的代碼結構

TensorFlow的自由度挺高的，類比的話，有那么點DL框架里的C++的意思；Python的語言靈活度也是高得不行，兩個很靈活的玩意放一起，寫個簡單模型自然想怎么玩就怎么玩，寫個復雜一些的模型，為了保證寫著方便，用著方便，改起來方便，還是需要比較好的代碼結構的。

如果翻翻GitHub上一些比較熱的用TF寫的模型，通常都會發現大家比較習慣于把代碼分成op、module和model三個部分。

op里是一些通用層或者運算的簡化定義，例如寫個卷積層，總是包含定義變量和定義運算。習慣于Keras這樣不需要自己定義變量的玩意當然不會太糾結，但用TF時，若是寫兩行定義一下變量總是挺讓人傷神的。

如果參照Keras的實現，通過寫個類來定制op，變量管理看起來方便一點，未免太過繁瑣。實際上TF提供的variable scope已經非常方便了，這一部分寫成這樣似乎也不錯

def conv2d(input, filter, kernel, strides=1, stddev=0.02, name='conv2d'):
   with tf.variable_scope(name):
       w = tf.get_variable(
           'w',
           (kernel, kernel, input.get_shape()[-1], filter),
           initializer=tf.truncated_normal_initializer(stddev=stddev)
       )
       conv = tf.nn.conv2d(input, w, strides=[1, strides, strides, 1], padding='VALID')
       b = tf.get_variable(
           'b',
           [filter],
           initializer=tf.constant_initializer(0.0)
       )
       conv = tf.reshape(tf.nn.bias_add(conv, b), tf.shape(conv))
       return conv

這樣定義幾個op之后，寫起代碼來就更有點類似于mxnet那樣的感覺了。

特別的，有些時候有些簡單結構，例如ResNet中的一個block這樣的玩意，我們也可以用類似的方式，用一個簡單函數包裝起來

def res_block(x, dim, name='res_block'):
   with tf.variable_scope(name):
       y = reflect_pad(x, name='rp1')
       y = conv2d(y, dim, 3, name='conv1')
       y = lrelu(y)
       y = reflect_pad(y, name='rp2')
       y = conv2d(y, dim, 3, name='conv2')
       y = lrelu(y)
       return tf.add(x, y)

對于重復的模塊，這樣的包裝也方便多次使用。

這些是很常見的做法。同時我們也發現了，幾乎每個這樣的函數里都少不了一個variable scope的使用，一方面避免定義變量時名字的重復以及訓練時變量的管理，另一方面也方便TensorBoard畫圖的時候能把有用的東西放到一起。但這樣每個函數里帶個name參數的做法寫多了也會煩，加上奇怪的縮進……我會更傾向于用一個裝飾器來解決這樣的問題，同時也能減少『忘了用variable scope』的情況。

def scope(default_name):
   def deco(fn):
       def wrapper(*args, **kwargs):
           if 'name' in kwargs:
               name = kwargs['name']
               kwargs.pop('name')
           else:
               name = default_name
           with tf.variable_scope(name):
               return fn(*args, **kwargs)
       return wrapper
   return deco@scope('conv2d')def conv2d(input, filter, kernel, strides=1, stddev=0.02):
   w = tf.get_variable(
       'w',
       (kernel, kernel, input.get_shape()[-1], filter),
       initializer=tf.truncated_normal_initializer(stddev=stddev)
   )
   conv = tf.nn.conv2d(input, w, strides=[1, strides, strides, 1], padding='VALID')
   b = tf.get_variable(
       'b',
       [filter],
       initializer=tf.constant_initializer(0.0)
   )
   conv = tf.reshape(tf.nn.bias_add(conv, b), tf.shape(conv))
   return conv

至于module，也就是一些稍微復雜的成型結構，例如GAN里的Discriminator和Generator，講道理這玩意其實和op大體上是類似的，就不多說了。

最后是model。通常大家都是用類來做，因為model中往往還包含了輸入數據用的placeholder、訓練用的op，甚至一些具體的方法等等內容。這一塊的代碼建議，只不過是最好先寫一個抽象類，把需要的幾個接口給定義一下，然后讓實際的model類繼承，代碼會漂亮很多，也更便于利用諸如PyCharm這樣的IDE來提示你哪些東西該做而沒有做。

關于config/options

網上常見的代碼里，模型的一些參數信息大都設計成用命令行參數來傳入，更多是直接使用tf.flags來處理。但無論如何，我仍然覺得定義一個config類來管理參數是有一定必要性的，直接使用tf.flags主要是是有大段tf.flags.DEFINE_xxx，不好看，也不方便直觀地反應默認參數。相對的，如果定義一個參數類，在__init__里寫下默認參數，然后寫個小方法自動地根據dir來添加這些tf.flags會漂亮許多。但這個只是個人觀點，似乎并沒有具體的優劣之分。

關于TensorBoard

不得不說TensorBoard作為TF自帶的配套可視化工具，只要你不是太在意刷新頻率的問題（通常不會有人在意這個吧……），用起來實在太方便。加上能夠自動生成運算的各個符號的結構圖，哪怕不說訓練，就是檢查模型結構是否符合自己所想都是個非常好用的工具。比如封面圖，生成出來用來檢查代碼的模型邏輯，還可以根據需要點選觀察依賴關系。

順帶一提，如果生成的模型圖長得非常奇怪，八成是代碼有問題……

不過要用好TensorBoard，有幾個小小的要點：首先是，至少，你的各個op和module里，得用上variable scope或者name scope。對于一個scope，在TensorBoard的Graph里會將其聚集成一個小塊，內部結構可以展開觀察，而如果不用scope，你會看到滿眼都是一堆一堆的基本op，當模型復雜時，圖基本沒法看……

此外，對于圖片處理，用好TensorBoard的ImageSummary當然是很不錯的選擇。但是記得一定要為添加圖片的summary op定義一個喂數據的placeholder。

self.p_img = tf.placeholder(tf.float32, shape=[1, 256 * 6, 256 * 4, 3])

self.img_op = tf.summary.image('sample', self.p_img)

……

img = np.array([img])

s_img = self.sess.run(self.img_op, feed_dict={self.p_img: img})

self.writer.add_summary(s_img, count)

這樣才是正確的。網上有些材料里告訴你可以直接用tf.summary.image('tag', data)來生成圖片summary，這樣其實每次都會構造一個新的summary，不便于圖片歸類，但更大的問題是，這樣做會使得每次都申請一個新的變量（用來裝你的圖片數據），倘若你有定周期存儲訓練權重的習慣，會發現沒幾個小時就會因為權重變量總量超過2GB而使得程序跑崩……想想看晚上跑著訓練的代碼想著可以回家休息了，結果前腳剛進家門，程序就罷工了，大好的訓練時間就給直接浪費了。

另外，這里的圖片可以是重新歸為0~255的整形的數據，也可以直接給浮點數據[-1, 1]。更不錯的想法是，先使用matplotlib/pil/numpy來合成、拼湊甚至生成圖像，然后再來添加，會讓效果更令人滿意，比如這樣：

最后補充一句……雙顯示器確實有利于提高寫代碼、改代碼以及碼字的效率……

雷峰網版權文章，未經授權禁止轉載。詳情見轉載須知。