作者丨蘇劍林

單位丨廣州火焰信息科技有限公司

研究方向丨NLP，神經網絡

個人主頁丨kexue.fm

今天我們來看一個小眾需求：自定義優化器。

細想之下，不管用什么框架，自定義優化器這個需求可謂真的是小眾中的小眾。一般而言，對于大多數任務我們都可以無腦地直接上 Adam，而調參煉丹高手一般會用 SGD 來調出更好的效果，換言之不管是高手新手，都很少會有自定義優化器的需求。

那這篇文章還有什么價值呢？有些場景下會有一點點作用。比如通過學習 Keras 中的優化器寫法，你可以對梯度下降等算法有進一步的認識，你還可以順帶看到 Keras 的源碼是多么簡潔優雅。

此外，有時候我們可以通過自定義優化器來實現自己的一些功能，比如給一些簡單的模型（例如 Word2Vec）重寫優化器（直接寫死梯度，而不是用自動求導），可以使得算法更快；自定義優化器還可以實現諸如“軟 batch”的功能。

Keras優化器

我們首先來看 Keras 中自帶優化器的代碼，位于：

https://github.com/keras-team/keras/blob/master/keras/optimizers.py

簡單起見，我們可以先挑 SGD 來看。當然，Keras 中的 SGD 算法已經把 momentum、nesterov、decay 等整合進去了，這使用起來方便，但不利于學習。所以我稍微簡化了一下，給出一個純粹的 SGD 算法的例子：?

from?keras.legacy?import?interfaces
from?keras.optimizers?import?Optimizer
from?keras?import?backend?as?K


class?SGD(Optimizer):
????"""Keras中簡單自定義SGD優化器
????"""

????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):
????????"""主要的參數更新算法
????????"""
????????grads?=?self.get_gradients(loss,?params)?#?獲取梯度
????????self.updates?=?[K.update_add(self.iterations,?1)]?#?定義賦值算子集合
????????self.weights?=?[self.iterations]?#?優化器帶來的權重，在保存模型時會被保存
????????for?p,?g?in?zip(params,?grads):
????????????#?梯度下降
????????????new_p?=?p?-?self.lr?*?g
????????????#?如果有約束，對參數加上約束
????????????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()))

應該不是解釋了吧？有沒有特別簡單的感覺？定義一個優化器也不是特別高大上的事情。

實現“軟batch”?

現在來實現一個稍微復雜一點的功能，就是所謂的“軟 batch”，不過我不大清楚是不是就叫這個名字，姑且先這樣叫著吧。大概的場景是：假如模型比較龐大，自己的顯卡最多也就能跑 batch size=16，但我又想起到 batch size=64 的效果，那可以怎么辦呢？

一種可以考慮的方案是，每次算 batch size=16，然后把梯度緩存起來，4 個 batch 后才更新參數。也就是說，每個小 batch 都算梯度，但每 4 個 batch 才更新一次參數。?

class?MySGD(Optimizer):
????"""Keras中簡單自定義SGD優化器
????每隔一定的batch才更新一次參數
????"""
????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):
????????"""主要的參數更新算法
????????"""
????????shapes?=?[K.int_shape(p)?for?p?in?params]
????????sum_grads?=?[K.zeros(shape)?for?shape?in?shapes]?#?平均梯度，用來梯度下降
????????grads?=?self.get_gradients(loss,?params)?#?當前batch梯度
????????self.updates?=?[K.update_add(self.iterations,?1)]?#?定義賦值算子集合
????????self.weights?=?[self.iterations]?+?sum_grads?#?優化器帶來的權重，在保存模型時會被保存
????????for?p,?g,?sg?in?zip(params,?grads,?sum_grads):
????????????#?梯度下降
????????????new_p?=?p?-?self.lr?*?sg?/?float(self.steps_per_update)
????????????#?如果有約束，對參數加上約束
????????????if?getattr(p,?'constraint',?None)?is?not?None:
????????????????new_p?=?p.constraint(new_p)
????????????cond?=?K.equal(self.iterations?%?self.steps_per_update,?0)
????????????#?滿足條件才更新參數
????????????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()))

應該也很容易理解吧。如果帶有動量的情況，寫起來復雜一點，但也是一樣的。重點就是引入多一個變量來儲存累積梯度，然后引入 cond 來控制是否更新，原來優化器要做的事情，都要在 cond 為 True 的情況下才做（梯度改為累積起來的梯度）。對比原始的 SGD，改動并不大。

“侵入式”優化器

上面實現優化器的方案是標準的，也就是按 Keras 的設計規范來做的，所以做起來很輕松。然而我曾經想要實現的一個優化器，卻不能用這種方式來實現，經過閱讀源碼，得到了一種“侵入式”的寫法，這種寫法類似“外掛”的形式，可以實現我需要的功能，但不是標準的寫法，在此也跟大家分享一下。

原始需求來源于之前的文章從動力學角度看優化算法SGD：一些小啟示，里邊指出梯度下降優化器可以看成是微分方程組的歐拉解法，進一步可以聯想到，微分方程組有很多比歐拉解法更高級的解法呀，能不能用到深度學習中？比如稍微高級一點的有“Heun 方法 [1]”：

其中 p 是參數（向量），g 是梯度，pi 表示 p 的第 i 次迭代時的結果。這個算法需要走兩步，大概意思就是普通的梯度下降先走一步（探路），然后根據探路的結果取平均，得到更精準的步伐，等價地可以改寫為：

這樣就清楚顯示出后面這一步實際上是對梯度下降的微調。?

但是實現這類算法卻有個難題，要計算兩次梯度，一次對參數 g(pi)，另一次對參數 p?i+1。而前面的優化器定義中 get_updates 這個方法卻只能執行一步（對應到 tf 框架中，就是執行一步 sess.run，熟悉 tf 的朋友知道單單執行一步 sess.run 很難實現這個需求），因此實現不了這種算法。

經過研究 Keras 模型的訓練源碼，我發現可以這樣寫：

class?HeunOptimizer:
????"""自定義Keras的侵入式優化器
????"""

????def?__init__(self,?lr):
????????self.lr?=?lr

????def?__call__(self,?model):
????????"""需要傳入模型，直接修改模型的訓練函數，而不按常規流程使用優化器，所以稱為“侵入式”
????????其實下面的大部分代碼，都是直接抄自keras的源碼：
????????https://github.com/keras-team/keras/blob/master/keras/engine/training.py#L491
????????也就是keras中的_make_train_function函數。
????????"""
????????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封裝為一個單獨的函數
????????????????????#?一個K.function就是一次sess.run
????????????????????for?f?in?update_functions:
????????????????????????_?=?f(ins)
????????????????????return?_

????????????????#?最后只需要將model的train_function屬性改為對應的函數
????????????????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)

其中關鍵思想在代碼中已經注釋了，主要是 Keras 的優化器最終都會被包裝為一個 train_function，所以我們只需要參照 Keras 的源碼設計好 train_function，并在其中插入我們自己的操作。在這個過程中，需要留意到 K.function 所定義的操作相當于一次 sess.run 就行了。

注：類似地還可以實現 RK23、RK45 等算法。遺憾的是，這種優化器缺很容易過擬合，也就是很容易將訓練集的 loss 降到很低，但是驗證集的 loss 和準確率都很差。

優雅的Keras

本文講了一個非常非常小眾的需求：自定義優化器，介紹了一般情況下 Keras 優化器的寫法，以及一種“侵入式”的寫法。如果真有這么個特殊需求，可以參考使用。

通過 Keras 中優化器的分析研究，我們進一步可以觀察到 Keras 整體代碼實在是非常簡潔優雅，難以挑剔。

參考文獻

[1].?https://en.wikipedia.org/wiki/Heun%27s_method

點擊以下標題查看作者其他文章：?

• 從無監督構建詞庫看「最小熵原理」

• 基于CNN的閱讀理解式問答模型：DGCNN

• 再談最小熵原理：飛象過河之句模版和語言結構

• 再談變分自編碼器VAE：從貝葉斯觀點出發
  

• 變分自編碼器VAE：這樣做為什么能成？

• 全新視角：用變分推斷統一理解生成模型

• 細水長flow之NICE：流模型的基本概念與實現
  

#投 稿 通 道#

?讓你的論文被更多人看到?

如何才能讓更多的優質內容以更短路徑到達讀者群體，縮短讀者尋找優質內容的成本呢？?答案就是：你不認識的人。

總有一些你不認識的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋梁，促使不同背景、不同方向的學者和學術靈感相互碰撞，迸發出更多的可能性。?

PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優質內容，可以是最新論文解讀，也可以是學習心得或技術干貨。我們的目的只有一個，讓知識真正流動起來。

總結

以上是生活随笔為你收集整理的Keras之小众需求：自定义优化器的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。