從一次CycleGAN實現聊聊TF

老婆大人的文已發布，地址在CycleGAN（以及DiscoGAN和DualGAN）簡介 - 知乎專欄

——————————————————

我還記得那天天氣挺熱，妻子還在為自己的研究方向發愁，然後就在同一天，我們倆都看到了同一篇文章，Unpaired Image-to-Image Translationnusing Cycle-Consistent Adversarial Networks [https://arxiv.org/pdf/1703.10593.pdf]

CycleGAN是個很有趣的想法，看完之後我倆都覺得非常興奮，並且隱隱地覺得，這後面有更多的內容可以挖。兩周後，妻子就這篇文章作了她研究生階段的第一次組內發表，而我也盡我所能做出了各種嘗試，努力發掘更多的可能性。

實現過程中可以說還是略微糾結的，最初是用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):n with tf.variable_scope(name):n w = tf.get_variable(n w,n (kernel, kernel, input.get_shape()[-1], filter),n initializer=tf.truncated_normal_initializer(stddev=stddev)n )n conv = tf.nn.conv2d(input, w, strides=[1, strides, strides, 1], padding=VALID)n b = tf.get_variable(n b,n [filter],n initializer=tf.constant_initializer(0.0)n )n conv = tf.reshape(tf.nn.bias_add(conv, b), tf.shape(conv))n return convn

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

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

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

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

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

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

至於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])nself.img_op = tf.summary.image(sample, self.p_img)n……nimg = np.array([img])ns_img = self.sess.run(self.img_op, feed_dict={self.p_img: img})nself.writer.add_summary(s_img, count)n

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

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

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

推薦閱讀：