作者丨蘇劍林
單位丨廣州火焰信息科技有限公司
研究方向丨NLP，神經(jīng)網(wǎng)絡(luò)
個人主頁丨kexue.fm
今天我們來看一個小眾需求：自定義優(yōu)化器。
細想之下，不管用什么框架，自定義優(yōu)化器這個需求可謂真的是小眾中的小眾。一般而言，對于大多數(shù)任務(wù)我們都可以無腦地直接上 Adam，而調(diào)參煉丹高手一般會用 SGD 來調(diào)出更好的效果，換言之不管是高手新手，都很少會有自定義優(yōu)化器的需求。
那這篇文章還有什么價值呢？有些場景下會有一點點作用。比如通過學(xué)習(xí) Keras 中的優(yōu)化器寫法，你可以對梯度下降等算法有進一步的認(rèn)識，你還可以順帶看到 Keras 的源碼是多么簡潔優(yōu)雅。
此外，有時候我們可以通過自定義優(yōu)化器來實現(xiàn)自己的一些功能，比如給一些簡單的模型（例如 Word2Vec）重寫優(yōu)化器（直接寫死梯度，而不是用自動求導(dǎo)），可以使得算法更快；自定義優(yōu)化器還可以實現(xiàn)諸如“軟 batch”的功能。
Keras優(yōu)化器
我們首先來看 Keras 中自帶優(yōu)化器的代碼，位于：
https://github.com/keras-team/keras/blob/master/keras/optimizers.py
簡單起見，我們可以先挑 SGD 來看。當(dāng)然，Keras 中的 SGD 算法已經(jīng)把 momentum、nesterov、decay 等整合進去了，這使用起來方便，但不利于學(xué)習(xí)。所以我稍微簡化了一下，給出一個純粹的 SGD 算法的例子：
from keras.legacy import interfaces
from keras.optimizers import Optimizer
from keras import backend as K
class SGD(Optimizer):
'''Keras中簡單自定義SGD優(yōu)化器
'''
def __init__(self, lr=0.01, **kwargs):
super(SGD, self).__init__(**kwargs)
with K.name_scope(self.__class__.__name__):
self.iterations = K.variable(0, dtype='int64', name='iterations')
self.lr = K.variable(lr, name='lr')
@interfaces.legacy_get_updates_support
def get_updates(self, loss, params):
'''主要的參數(shù)更新算法
'''
grads = self.get_gradients(loss, params) # 獲取梯度
self.updates = [K.update_add(self.iterations, 1)] # 定義賦值算子集合
self.weights = [self.iterations] # 優(yōu)化器帶來的權(quán)重，在保存模型時會被保存
for p, g in zip(params, grads):
# 梯度下降
new_p = p - self.lr * g
# 如果有約束，對參數(shù)加上約束
if getattr(p, 'constraint', None) is not None:
new_p = p.constraint(new_p)
# 添加賦值
self.updates.append(K.update(p, new_p))
return self.updates
def get_config(self):
config = {'lr': float(K.get_value(self.lr))}
base_config = super(SGD, self).get_config()
return dict(list(base_config.items()) + list(config.items()))
應(yīng)該不是解釋了吧？有沒有特別簡單的感覺？定義一個優(yōu)化器也不是特別高大上的事情。
實現(xiàn)“軟batch”
現(xiàn)在來實現(xiàn)一個稍微復(fù)雜一點的功能，就是所謂的“軟 batch”，不過我不大清楚是不是就叫這個名字，姑且先這樣叫著吧。大概的場景是：假如模型比較龐大，自己的顯卡最多也就能跑 batch size=16，但我又想起到 batch size=64 的效果，那可以怎么辦呢？
一種可以考慮的方案是，每次算 batch size=16，然后把梯度緩存起來，4 個 batch 后才更新參數(shù)。也就是說，每個小 batch 都算梯度，但每 4 個 batch 才更新一次參數(shù)。
class MySGD(Optimizer):
'''Keras中簡單自定義SGD優(yōu)化器
每隔一定的batch才更新一次參數(shù)
'''
def __init__(self, lr=0.01, steps_per_update=1, **kwargs):
super(MySGD, self).__init__(**kwargs)
with K.name_scope(self.__class__.__name__):
self.iterations = K.variable(0, dtype='int64', name='iterations')
self.lr = K.variable(lr, name='lr')
self.steps_per_update = steps_per_update # 多少batch才更新一次
@interfaces.legacy_get_updates_support
def get_updates(self, loss, params):
'''主要的參數(shù)更新算法
'''
shapes = [K.int_shape(p) for p in params]
sum_grads = [K.zeros(shape) for shape in shapes] # 平均梯度，用來梯度下降
grads = self.get_gradients(loss, params) # 當(dāng)前batch梯度
self.updates = [K.update_add(self.iterations, 1)] # 定義賦值算子集合
self.weights = [self.iterations] + sum_grads # 優(yōu)化器帶來的權(quán)重，在保存模型時會被保存
for p, g, sg in zip(params, grads, sum_grads):
# 梯度下降
new_p = p - self.lr * sg / float(self.steps_per_update)
# 如果有約束，對參數(shù)加上約束
if getattr(p, 'constraint', None) is not None:
new_p = p.constraint(new_p)
cond = K.equal(self.iterations % self.steps_per_update, 0)
# 滿足條件才更新參數(shù)
self.updates.append(K.switch(cond, K.update(p, new_p), p))
# 滿足條件就要重新累積，不滿足條件直接累積
self.updates.append(K.switch(cond, K.update(sg, g), K.update(sg, sg+g)))
return self.updates
def get_config(self):
config = {'lr': float(K.get_value(self.lr)),
'steps_per_update': self.steps_per_update}
base_config = super(MySGD, self).get_config()
return dict(list(base_config.items()) + list(config.items()))
應(yīng)該也很容易理解吧。如果帶有動量的情況，寫起來復(fù)雜一點，但也是一樣的。重點就是引入多一個變量來儲存累積梯度，然后引入 cond 來控制是否更新，原來優(yōu)化器要做的事情，都要在 cond 為 True 的情況下才做（梯度改為累積起來的梯度）。對比原始的 SGD，改動并不大。
“侵入式”優(yōu)化器
上面實現(xiàn)優(yōu)化器的方案是標(biāo)準(zhǔn)的，也就是按 Keras 的設(shè)計規(guī)范來做的，所以做起來很輕松。然而我曾經(jīng)想要實現(xiàn)的一個優(yōu)化器，卻不能用這種方式來實現(xiàn)，經(jīng)過閱讀源碼，得到了一種“侵入式”的寫法，這種寫法類似“外掛”的形式，可以實現(xiàn)我需要的功能，但不是標(biāo)準(zhǔn)的寫法，在此也跟大家分享一下。
原始需求來源于之前的文章從動力學(xué)角度看優(yōu)化算法SGD：一些小啟示，里邊指出梯度下降優(yōu)化器可以看成是微分方程組的歐拉解法，進一步可以聯(lián)想到，微分方程組有很多比歐拉解法更高級的解法呀，能不能用到深度學(xué)習(xí)中？比如稍微高級一點的有“Heun 方法 [1]”：
其中 p 是參數(shù)（向量），g 是梯度，pi 表示 p 的第 i 次迭代時的結(jié)果。這個算法需要走兩步，大概意思就是普通的梯度下降先走一步（探路），然后根據(jù)探路的結(jié)果取平均，得到更精準(zhǔn)的步伐，等價地可以改寫為：
這樣就清楚顯示出后面這一步實際上是對梯度下降的微調(diào)。
但是實現(xiàn)這類算法卻有個難題，要計算兩次梯度，一次對參數(shù) g(pi)，另一次對參數(shù) p?i+1。而前面的優(yōu)化器定義中 get_updates 這個方法卻只能執(zhí)行一步（對應(yīng)到 tf 框架中，就是執(zhí)行一步 sess.run，熟悉 tf 的朋友知道單單執(zhí)行一步 sess.run 很難實現(xiàn)這個需求），因此實現(xiàn)不了這種算法。
經(jīng)過研究 Keras 模型的訓(xùn)練源碼，我發(fā)現(xiàn)可以這樣寫：
class HeunOptimizer:
'''自定義Keras的侵入式優(yōu)化器
'''
def __init__(self, lr):
self.lr = lr
def __call__(self, model):
'''需要傳入模型，直接修改模型的訓(xùn)練函數(shù)，而不按常規(guī)流程使用優(yōu)化器，所以稱為“侵入式”
其實下面的大部分代碼，都是直接抄自keras的源碼：
https://github.com/keras-team/keras/blob/master/keras/engine/training.py#L491
也就是keras中的_make_train_function函數(shù)。
'''
params = model._collected_trainable_weights
loss = model.total_loss
inputs = (model._feed_inputs +
model._feed_targets +
model._feed_sample_weights)
inputs += [K.learning_phase()]
with K.name_scope('training'):
with K.name_scope('heun_optimizer'):
old_grads = [[K.zeros(K.int_shape(p)) for p in params]]
update_functions = []
for i,step in enumerate([self.step1, self.step2]):
updates = (model.updates +
step(loss, params, old_grads) +
model.metrics_updates)
# 給每一步定義一個K.function
updates = K.function(inputs,
[model.total_loss] + model.metrics_tensors,
updates=updates,
name='train_function_%s'%i,
**model._function_kwargs)
update_functions.append(updates)
def F(ins):
# 將多個K.function封裝為一個單獨的函數(shù)
# 一個K.function就是一次sess.run
for f in update_functions:
_ = f(ins)
return _
# 最后只需要將model的train_function屬性改為對應(yīng)的函數(shù)
model.train_function = F
def step1(self, loss, params, old_grads):
ops = []
grads = K.gradients(loss, params)
for p,g,og in zip(params, grads, old_grads[0]):
ops.append(K.update(og, g))
ops.append(K.update(p, p - self.lr * g))
return ops
def step2(self, loss, params, old_grads):
ops = []
grads = K.gradients(loss, params)
for p,g,og in zip(params, grads, old_grads[0]):
ops.append(K.update(p, p - 0.5 * self.lr * (g - og)))
return ops
用法是：
opt = HeunOptimizer(0.1)
opt(model)
model.fit(x_train, y_train, epochs=100, batch_size=32)
其中關(guān)鍵思想在代碼中已經(jīng)注釋了，主要是 Keras 的優(yōu)化器最終都會被包裝為一個 train_function，所以我們只需要參照 Keras 的源碼設(shè)計好 train_function，并在其中插入我們自己的操作。在這個過程中，需要留意到 K.function 所定義的操作相當(dāng)于一次 sess.run 就行了。
注：類似地還可以實現(xiàn) RK23、RK45 等算法。遺憾的是，這種優(yōu)化器缺很容易過擬合，也就是很容易將訓(xùn)練集的 loss 降到很低，但是驗證集的 loss 和準(zhǔn)確率都很差。
優(yōu)雅的Keras
本文講了一個非常非常小眾的需求：自定義優(yōu)化器，介紹了一般情況下 Keras 優(yōu)化器的寫法，以及一種“侵入式”的寫法。如果真有這么個特殊需求，可以參考使用。
通過 Keras 中優(yōu)化器的分析研究，我們進一步可以觀察到 Keras 整體代碼實在是非常簡潔優(yōu)雅，難以挑剔。
參考文獻
[1]. https://en./wiki/Heun%27s_method
