“让Keras更酷一些!”:小众的自定义优化器
By 苏剑林 | 2018-09-08 | 86017位读者 |沿着之前的《“让Keras更酷一些!”:精巧的层与花式的回调》写下去~
今天我们来看一个小众需求:自定义优化器。
细想之下,不管用什么框架,自定义优化器这个需求可谓真的是小众中的小众。一般而言,对于大多数任务我们都可以无脑地直接上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才更新一次参数。
2019年7月8日更新:以下实现是有误的,无法达到预期效果,修正的实现请看《Keras梯度累积优化器:用时间换取效果》。
如果真的有这个需求,那么就只能通过修改优化器来解决了。在前面的SGD的基础上,参考代码如下:
class7 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方法”:
$$\begin{aligned}\tilde{p}_{i+1} =& p_i + \epsilon g(p_i)\\
p_{i+1} =& p_i + \frac{1}{2}\epsilon \big[g(p_i)+g(\tilde{p}_{i+1})\big]
\end{aligned}$$
其中$p$是参数(向量),$g$是梯度,$p_i$表示$p$的第$i$次迭代时的结果。这个算法需要走两步,大概意思就是普通的梯度下降先走一步(探路),然后根据探路的结果取平均,得到更精准的步伐,等价地可以改写为:
$$\begin{aligned}\tilde{p}_{i+1} =& p_i + \epsilon g(p_i)\\
p_{i+1} =& \tilde{p}_{i+1} + \frac{1}{2}\epsilon \big[g(\tilde{p}_{i+1}) - g(p_i)\big]
\end{aligned}$$
这样就清楚显示出后面这一步实际上是对梯度下降的微调。
但是实现这类算法却有个难题,要计算两次梯度,一次对参数$p_i$,另一次对参数$\tilde{p}_{i+1}$。而前面的优化器定义中get_updates这个方法却只能执行一步(对应到tf框架中,就是执行一步sess.run,熟悉tf的朋友知道单单执行一步sess.run很难实现这个需求),因此实现不了这种算法。经过研究Keras模型的训练源码,我发现可以这样写:
#! -*- coding: utf-8 -*-
from keras.optimizers import Optimizer
from keras import backend as K
class InjectOptimizer(Optimizer):
"""定义注入式优化器的基类
需要传入模型,直接修改模型的训练函数,而不按常规流程使用优化器,所以称为“侵入式”
其实下面的大部分代码,都是直接抄自keras的源码:
https://github.com/keras-team/keras/blob/master/keras/engine/training.py#L497
也就是keras中的_make_train_function函数。
"""
def get_updates(self, loss, params):
return []
def get_grouped_updates(self, loss, params):
raise NotImplementedError
def inject(self, model):
"""传入模型做注入
"""
if not hasattr(model, 'train_function'):
raise RuntimeError('You must compile your model before using it.')
model._check_trainable_weights_consistency()
if model.train_function is None:
inputs = (model._feed_inputs +
model._feed_targets +
model._feed_sample_weights)
if model._uses_dynamic_learning_phase():
inputs += [K.learning_phase()]
with K.name_scope('training'):
train_functions = []
with K.name_scope(model.optimizer.__class__.__name__):
grouped_training_updates = self.get_grouped_updates(
params=model._collected_trainable_weights,
loss=model.total_loss)
for i, updates in enumerate(grouped_training_updates[1:]):
f = K.function(
inputs,
[model.total_loss],
updates=updates,
name='train_function_%s' % (i + 1),
**model._function_kwargs)
train_functions.append(f)
# Gets loss and metrics. Updates weights at each call.
first_updates = (model.updates +
grouped_training_updates[0],
model.metrics_updates)
first_train = K.function(
inputs,
[model.total_loss] + model.metrics_tensors,
updates=first_updates,
name='train_function',
**model._function_kwargs)
def F(inputs):
R = first_train(inputs)
for f in train_functions:
f(inputs)
return R
model.train_function = F
class HeunOptimizer(InjectOptimizer):
"""Heun优化器
( https://en.wikipedia.org/wiki/Heun%27s_method )
"""
def __init__(self, lr, **kwargs):
super(HeunOptimizer, self).__init__(**kwargs)
with K.name_scope(self.__class__.__name__):
self.lr = K.variable(lr, name='lr')
def get_updates_1(self, loss, params, cache_grads):
updates = []
grads = self.get_gradients(loss, params)
for p, g, cg in zip(params, grads, cache_grads):
updates.append(K.update(cg, g))
updates.append(K.update(p, p - self.lr * g))
return updates
def get_updates_2(self, loss, params, cache_grads):
updates = []
grads = self.get_gradients(loss, params)
for p, g, cg in zip(params, grads, cache_grads):
updates.append(K.update(p, p - 0.5 * self.lr * (g - cg)))
return updates
def get_grouped_updates(self, loss, params):
cache_grads = [K.zeros(K.int_shape(p)) for p in params]
return [
self.get_updates_1(loss, params, cache_grads),
self.get_updates_2(loss, params, cache_grads)
]
用法是:
opt = HeunOptimizer(0.1)
model.compile(loss='mse', optimizer=opt)
opt.inject(model) # 必须执行这步
model.fit(x_train, y_train, epochs=100, batch_size=32)
其中关键思想在代码中已经注释了,主要是Keras的优化器最终都会被包装为一个train_function,所以我们只需要参照Keras的源码设计好train_function,并在其中插入我们自己的操作。在这个过程中,需要留意到K.function所定义的操作相当于一次sess.run就行了。
注:类似地还可以实现RK23、RK45等算法。遗憾的是,这种优化器效果并不好~
优雅的Keras #
本文讲了一个非常非常小众的需求:自定义优化器,介绍了一般情况下Keras优化器的写法,以及一种“侵入式”的写法。如果真有这么个特殊需求,可以参考使用。
通过Keras中优化器的分析研究,我们进一步可以观察到Keras整体代码实在是非常简洁优雅,难以挑剔~
转载到请包括本文地址:https://kexue.fm/archives/5879
更详细的转载事宜请参考:《科学空间FAQ》
如果您还有什么疑惑或建议,欢迎在下方评论区继续讨论。
如果您觉得本文还不错,欢迎分享/打赏本文。打赏并非要从中获得收益,而是希望知道科学空间获得了多少读者的真心关注。当然,如果你无视它,也不会影响你的阅读。再次表示欢迎和感谢!
如果您需要引用本文,请参考:
苏剑林. (Sep. 08, 2018). 《“让Keras更酷一些!”:小众的自定义优化器 》[Blog post]. Retrieved from https://kexue.fm/archives/5879
@online{kexuefm-5879,
title={“让Keras更酷一些!”:小众的自定义优化器},
author={苏剑林},
year={2018},
month={Sep},
url={\url{https://kexue.fm/archives/5879}},
}
September 9th, 2018
博主,你不要停,我要一辈子你看的博客!加油!
September 12th, 2018
正需要,博主就更了~ 赞
September 16th, 2018
谢谢,很好的文章,学习了https://m.cnoz.org/0_15/
November 15th, 2018
博主你好,请问那个自定义Keras的侵入式优化器中class HeunOptimizer应该怎么使用啊?是把它嵌在optimizer.py文件里面吗?
已经说了怎么用了呀,定义好这个类,定义好模型,然后:
opt = HeunOptimizer(0.1)
opt(model)
April 28th, 2019
hello!最近我在用keras 做一个模型,我想把它放在手机上,不知道该怎么弄呢?可以简单一下说说吗?期待你的回复!谢谢!
October 29th, 2019
楼主你好,请问那个"软batch"优化器,我看你在累积梯度那里说是因为“K.switch(cond, K.update(p, new_p), p)”的原因,那么改为“K.update(K.switch(cond, new_p, p))”是不是就可以了呢?
写错了,应该是““K.update(p, K.switch(cond, new_p, p))”
这倒是可以~
December 4th, 2019
您好,我跑了一下那个侵入式代码,结果显示AttributeError: can't set attribute
经过查看,出错的是这两条语句,opt = HeunOptimizer(0.1) self.lr = K.variable(lr, name='lr')请问要怎么改才能不出错呢,不好意思,我是个菜鸟
用keras 2.2.4
参考您的回复,问题已经解决,谢谢您,请问有GitHub上的完整版代码吗?
没有,这个优化器目前看来并没有什么意思,不再维护。
February 10th, 2020
[...]In addition, it seems that the “hacking” trick of the repo comes from the following article (judging from his code and comments): https://kexue.fm/archives/5879/comment-page-1 (Sorry it is[...]
June 14th, 2020
我在使用新版keras的时候,会报错没有self.lr这个属性,这个是因为lr被改成了learning_rate,但我改成learning_rate之后,所有的optimizer都不work(包括您的例子),您有没有遇到这样的问题。
就是可以跑通,但没有优化的能力
找到问题了,需要在get_update()之前加上@K.symbolic。 另外非常佩服您,一个人发表了如此的多的内容,涵盖的知识面也非常广。
谢谢夸奖。
恭喜你成功找到了原因。keras 2.3.x兼容了tf 2.x,tf 2.x默认是动态图,而keras使用的是静态图,所以需要@K.symbolic来进入静态图模式。
个人还有一个疑问不知可否赐教,我想自己设计在多设备上训练的梯度融合规则,比如说在一次梯度融合中,使用梯度的加权平均,或者是想每5次梯度更新上传一次梯度共,其余的时候每个设备不上传梯度,请问这样的功能用keras需要在哪里实现,如果能指点一下tensorflow的实现也是极好的。非常感谢。
这个层面的内容我已经不懂了,抱歉~
June 19th, 2020
还是非常感谢,您的网站里keras相关的东西还有很多值得我学习的地方。