知識星球
作者丨蘇劍林
單位丨追一科技
研究方向丨NLP，神經網路
個人主頁丨kexue.fm
繼續“讓Keras更酷一些”之旅。
今天我們會用 Keras 實現靈活地輸出任意中間變數，還有無縫地進行權重滑動平均，最後順便介紹一下生成器的行程安全寫法。
首先是輸出中間變數。在自定義層時，我們可能希望檢視中間變數，這些需求有些是比較容易實現的，比如檢視中間某個層的輸出，只需要將截止到這個層的部分模型儲存為一個新模型即可，但有些需求是比較困難的，比如在使用 Attention 層時我們可能希望檢視那個 Attention 矩陣的值，如果用構建新模型的方法則會非常麻煩。而本文則給出一種簡單的方法，徹底滿足這個需求。
接著是權重滑動平均。權重滑動平均是穩定、加速模型訓練甚至提升模型效果的一種有效方法，很多大型模型（尤其是 GAN）幾乎都用到了權重滑動平均。一般來說權重滑動平均是作為最佳化器的一部分，所以一般需要重寫最佳化器才能實現它。本文介紹一個權重滑動平均的實現，它可以無縫插入到任意 Keras 模型中，不需要自定義最佳化器。
至於生成器的行程安全寫法，則是因為 Keras 讀取生成器的時候，用到了多行程，如果生成器本身也包含了一些多行程操作，那麼可能就會導致異常，所以需要解決這個這個問題。
輸出中間變數
這一節以基本模型為例，逐步深入地介紹如何獲取 Keras 的中間變數。
x_in = Input(shape=(784,))
x = x_in
x = Dense(512, activation='relu')(x)
x = Dropout(0.2)(x)
x = Dense(256, activation='relu')(x)
x = Dropout(0.2)(x)
x = Dense(num_classes, activation='softmax')(x)
model = Model(x_in, x)
作為一個新模型
假如模型訓練完成後，我想要獲取 x = Dense(256, activation=’relu’)(x) 對應的輸出，那可以在定義模型的時候，先把對應的變數存起來，然後重新定義一個模型：
x_in = Input(shape=(784,))
x = x_in
x = Dense(512, activation='relu')(x)
x = Dropout(0.2)(x)
x = Dense(256, activation='relu')(x)
y = x
x = Dropout(0.2)(x)
x = Dense(num_classes, activation='softmax')(x)
model = Model(x_in, x)
model2 = Model(x_in, y)
將 model 訓練完成後，直接用 model2.predict 就可以檢視對應的 256 維的輸出了。這樣做的前提是 y 必須是某個層的輸出，不能是隨意一個張量。
K.function！
有時候我們自定義了一個比較複雜的層，比較典型的就是 Attention 層，我們希望檢視層的一些中間變數，比如對應的 Attention 矩陣，這時候就比較麻煩了，如果想要用前面的方式，那麼就要把原來的 Attention 層分開為兩個層定義才行。
因為前面已經說了，新定義一個 Keras 模型時輸入輸出都必須是 Keras 層的輸入輸出，不能是隨意一個張量。這樣一來，如果想要分別檢視層的多個中間變數，那就要將層不斷地拆開為多個層來定義，顯然是不夠友好的。
其實 Keras 提供了一個終極的解決方案： K.function ！
介紹 K.function 之前，我們先寫一個簡單示例：
class Normal(Layer):
def __init__(self, **kwargs):
super(Normal, self).__init__(**kwargs)
def build(self, input_shape):
self.kernel = self.add_weight(name='kernel',
shape=(1,),
initializer='zeros',
trainable=True)
self.built = True
def call(self, x):
self.x_normalized = K.l2_normalize(x, -1)
return self.x_normalized * self.kernel
x_in = Input(shape=(784,))
x = x_in
x = Dense(512, activation='relu')(x)
x = Dropout(0.2)(x)
x = Dense(256, activation='relu')(x)
x = Dropout(0.2)(x)
normal = Normal()
x = normal(x)
x = Dense(num_classes, activation='softmax')(x)
model = Model(x_in, x)
在上面的例子中， Normal 定義了一個層，層的輸出是 self.x_normalized * self.kernel ，不過我想在訓練完成後獲取 self.x_normalized 的值，而它是跟輸入有關，並且不是一個層的輸出。這樣一來前面的方法就沒法用了，但用 K.function 就只是一行程式碼：
fn = K.function([x_in], [normal.x_normalized])
K.function 的用法跟定義一個新模型類似，要把所有跟 normal.x_normalized 相關的輸入張量都傳進去，但是不要求輸出是一個層的輸出，允許是任意張量！傳回的 fn 是一個具有函式功能的物件，所以只需要：
就可以獲取到 x_test 對應的 x_normalized 了！比定義一個新模型簡單通用多了。
事實上 K.function 就是 Keras 底層的基礎函式之一，它直接封裝好了後端的輸入輸出操作，換句話說，你用 Tensorflow 為後端時， fn([x_test]) 就相當於：
sess.run(normal.x_normalized, feed_dict={x_in: x_test})
所以 K.function 的輸出允許是任意張量，因為它本來就在直接操作後端了。
權重滑動平均
權重滑動平均是提供訓練穩定性的有效方法，透過滑動平均可以幾乎零額外成本地提高解的效能。權重滑動平均一般就是指“Exponential Moving Average”，簡稱 EMA，這是因為一般滑動平均時會使用指數衰減作為權重的比例。它已經被主流模型所接受，尤其是 GAN，在很多 GAN 論文中我們通常會看到類似的描述：
we use an exponential moving average with decay 0.999 over the weight …
這就意味著 GAN 模型使用了 EMA。此外，普通模型也會使用，比如 QANet: Combining Local Convolution with Global Self-Attention for Reading Comprehension 就在訓練過程中用了 EMA，衰減率是 0.9999。
滑動平均的格式
滑動平均的格式其實非常簡單，假設每次最佳化器的更新為：
這裡的 Δθn 就是最佳化器帶來的更新，最佳化器可以是 SGD、Adam 等任意一種。而滑動平均則是維護一組新的新的變數 Θ：
其中 α 是一個接近於 1 的正常數，稱為“衰減率（decay rate）”。
註意，儘管在形式上有點相似，但它跟動量加速不一樣：EMA 不改變原來最佳化器的軌跡，即原來最佳化器怎麼走，現在依然是同樣的走法，只不過它維護一組新變數，來平均原來最佳化器的軌跡；而動量加速則是改變了原來最佳化器的軌跡。
再次強調，權重滑動平均不改變最佳化器的走向，只不過它降最佳化器的最佳化軌跡上的點做了平均後，作為最終的模型權重。
關於權重滑動平均的原理和效果，可以進一步參考《從動力學角度看最佳化演演算法（四）：GAN 的第三個階段》一文。
巧妙的註入實現
實現 EMA 的要點是如何在原來最佳化器的基礎上引入一組新的平均變數，並且在每次引數更新後執行平均變數的更新。這需要對 Keras 的原始碼及其實現邏輯有一定的瞭解。
在此給出的參考實現如下：
class ExponentialMovingAverage:
"""對模型權重進行指數滑動平均。
用法：在model.compile之後、第一次訓練之前使用；
先初始化物件，然後執行inject方法。
"""
def __init__(self, model, momentum=0.9999):
self.momentum = momentum
self.model = model
self.ema_weights = [K.zeros(K.shape(w)) for w in model.weights]
def inject(self):
"""新增更新運算元到model.metrics_updates。
"""
self.initialize()
for w1, w2 in zip(self.ema_weights, self.model.weights):
op = K.moving_average_update(w1, w2, self.momentum)
self.model.metrics_updates.append(op)
def initialize(self):
"""ema_weights初始化跟原模型初始化一致。
"""
self.old_weights = K.batch_get_value(self.model.weights)
K.batch_set_value(zip(self.ema_weights, self.old_weights))
def apply_ema_weights(self):
"""備份原模型權重，然後將平均權重應用到模型上去。
"""
self.old_weights = K.batch_get_value(self.model.weights)
ema_weights = K.batch_get_value(self.ema_weights)
K.batch_set_value(zip(self.model.weights, ema_weights))
def reset_old_weights(self):
"""恢復模型到舊權重。
"""
K.batch_set_value(zip(self.model.weights, self.old_weights))
使用方法很簡單：
EMAer = ExponentialMovingAverage(model) # 在模型compile之後執行
EMAer.inject() # 在模型compile之後執行
model.fit(x_train, y_train) # 訓練模型
訓練完成後：
EMAer.apply_ema_weights() # 將EMA的權重應用到模型中
model.predict(x_test) # 進行預測、驗證、儲存等操作
EMAer.reset_old_weights() # 繼續訓練之前，要恢復模型舊權重。還是那句話，EMA不影響模型的最佳化軌跡。
model.fit(x_train, y_train) # 繼續訓練
現在翻看實現過程，可以發現主要的一點是引入了 K.moving_average_update 操作，並且插入到 model.metrics_updates 中，在訓練過程中，模型會讀取並執行 model.metrics_updates 的所有運算元，從而完成了滑動平均。
行程安全生成器
一般來說，當訓練資料無法全部載入記憶體，或者需要動態生成訓練資料時，就會用到 generator。一般來說，Keras 模型的 generator 的寫法是：
def data_generator():
while True:
x_train = something
y_train = otherthing
yield x_train, y_train
但如果 someting 或 otherthing 裡邊包含了多行程操作，就可能出問題。這時候有兩種解決方法，一是 fit_generator 時將設定引數use_multiprocessing=False, worker=0 ；另一種方法就是透過繼承 keras.utils.Sequence 類來寫生成器。
官方參考例子
官方對 keras.utils.Sequence 類的介紹如下：
https://keras.io/utils/#sequence
官方強調：
Sequence are a safer way to do multiprocessing. This structure guarantees that the network will only train once on each sample per epoch which is not the case with generators.
總之，就是對於多行程來說它是安全的，可以放心用。官方提供的例子如下：
from skimage.io import imread
from skimage.transform import resize
import numpy as np
# Here, `x_set` is list of path to the images
# and `y_set` are the associated classes.
class CIFAR10Sequence(Sequence):
def __init__(self, x_set, y_set, batch_size):
self.x, self.y = x_set, y_set
self.batch_size = batch_size
def __len__(self):
return int(np.ceil(len(self.x) / float(self.batch_size)))
def __getitem__(self, idx):
batch_x = self.x[idx * self.batch_size:(idx + 1) * self.batch_size]
batch_y = self.y[idx * self.batch_size:(idx + 1) * self.batch_size]
return np.array([
resize(imread(file_name), (200, 200))
for file_name in batch_x]), np.array(batch_y)
就是按格式定義好 __len__ 和 __getitem__ 方法就行了， __getitem__ 方法直接傳回一個 batch 的資料。
bert as service例子
我第一次發現 Sequence 的必要性，是在試驗 bert as service 的時候。bert as service 是肖涵大佬搞的一個快速獲取 bert 編碼向量的服務元件，我曾經想用它獲取字向量，然後傳入到 Keras 中訓練，但發現總會訓練著訓練著就卡住了。
經過搜尋，確認是 Keras 的 fit_generator 所帶的多行程，和 bert-as-service 自帶的多行程衝突問題，具體怎麼衝突我也比較模糊，就不深究了。而這裡提供了一個參考的解決方案，用的就是繼承 Sequence 類來寫生成器。
https://github.com/hanxiao/bert-as-service/issues/29#issuecomment-442362241
PS：就呼叫 bert as service 而言，後面肖涵大佬提供了協程版的 ConcurrentBertClient ，可以取代原來的 BertClient ，這樣哪怕在原始生成器也不會有問題了。
清流般的Keras
在我眼裡，Keras 就是深度學習框架中的一股清流，就好比 Python 是所有程式語言中的一股清流一樣。用 Keras 實現所需要做的事情，就好比一次次愜意的享受。
朋友會在“發現-看一看”看到你“在看”的內容