这两周投入了比较多的精力去做bert4keras的开发,除了一些API的规范化工作外,其余的主要工作量是构建预训练部分的代码。在昨天,预训练代码基本构建完毕,并同时在TPU/多GPU环境下测试通过,从而有志(有算力)改进预训练模型的同学多了一个选择。——这可能是目前最为清晰易懂的bert及其预训练代码。

预训练代码链接: https://github.com/bojone/bert4keras/tree/master/pretraining

经过这两周的开发(填坑),笔者的最大感想就是:Keras已经成为了tensorflow的黄金标准了。只要你的代码按照Keras的标准规范写,那可以轻松迁移到tf.keras中去,继而可以非常轻松地在TPU或多GPU环境下训练,真正的几乎是一劳永逸。相反,如果你的写法过于灵活,包括像笔者之前介绍的很多“移花接木”式的Keras技巧,就可能会有不少问题,甚至可能出现的一种情况是:就算你已经在多GPU上跑通了,在TPU上你也死活调不通。

Keras和Tensorflow

Keras和Tensorflow

不遗余力的支持 #

大家都说tensorflow 2.0主推tf.keras,但事实上,从tensorflow 1.14开始,tf.keras就已经成为了它的黄金标准了,所以如果大家要体验Google爸爸是如何不遗余力支持keras的,只需要tensorflow 1.14+就行了,不一定要升级到2.0。目前bert4keras的代码同时支持原版keras以及tf.keras,在通常的单卡finetune任务里,大家可以用keras或者tf.keras都行,但如果要用多GPU训练甚至用TPU训练,那最好的选择还是tf.keras。

要入门tf.keras,首先建议参考一个很不错的网站:https://tf.wiki/

在tf.keras中,将一个模型从单卡变成多卡训练,代码非常简单:

strategy = tf.distribute.MirroredStrategy()

with strategy.scope():
    model = create_a_model()
    model.compile(loss='mse', optimizer='adam')

model.fit(train_x, train_y, epochs=10)

也就是说,只要定义一个strategy,然后在这个strategyscope下建立的模型,就是多卡的模型了。多卡训练,从未如此简单~

顺便提一下,Keras本身自带了multi_gpu_model函数来实现多GPU训练,但笔者亲身测试multi_gpu_model并不好用,有时候还无法生效。总之还是推荐用tf.keras。另外,上面是单机多卡的写法,多机多卡类似,但我没有相应的环境测试,所以就没有给出例子,要测试的朋友,请参考https://tf.wiki/里边的介绍。

那TPU呢?一样简单,把strategy替换一下就行了(tensorflow 2.0有小改动):

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=tpu_address)
tf.config.experimental_connect_to_host(resolver.master())
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)

有没有感觉到简单得不可思议?这时候大家总该明白我前面说的“不遗余力”支持keras的含义了吧。从tensorflow 1.14开始,只要你用标准的keras写法,那么几乎可以无往而不利。

怎样才算是标准? #

前面不断强调用标准的keras写法,那怎样才算是标准呢?这里汇总一些经验。

1、尽可能全用keras自带的层、loss函数和优化器实现我们所需要的功能,如果一个模型全都是用keras内置的层、loss函数和优化器,那么基本上可以保证在多GPU或TPU上都能跑通了。

2、如果要自定义层,要严格按照规范来,尤其是要把get_config方法写好。测试自己写得规范与否的一个方法是:用你自定义的层去构建一个模型,然后看看能不能被clone_model函数成功克隆该模型,如果可以,说明你这个层的定义已经规范了。

3、如果要用TPU训练,不要在模型的最后用add_loss自定义loss,也不要用add_metric来添加metric。如果需要自定义复杂的loss或metric,请将它们定义为一个层的输出,参考这种写法

4、如果要用TPU训练,切忌在训练过程中使用动态(变长)的写法,比如使用tf.where时参数x,y不能为None,否则tf.where的结果长度不确定,还有tf中几乎所有带有dynamic字眼的函数都不能用。

可以看到,所谓的“标准”,其实就是最大程度上模仿着keras已有的写法了,尽量少自己创造。只要做到1、2点,就可以在tf.keras下轻易用多GPU训练了;后面3、4两点是针对TPU填的坑,总的来说就是一切要是静态的。

尽管tensorflow 2.0开始已经默认动态图了,但笔者并不推荐使用,个人认为我们应当去习惯静态图的模型构建流程。动态图虽然方便调试,但会让我们对这种即时的输出结果产生严重依赖,降低我们面对复杂问题时的debug能力。类似地,我也不推荐使用代码补全、代码提示等工具,这些工具会让我们产生过多依赖,使得我们不真正去了解你要用的函数本身。(个人观点,不喜勿喷。)

人性化的胜利 #

如果没有记错,笔者是2015年初接触的Keras,当时也没有太多深度学习框架,只是想找个趁手的工具实现几个简单的模型,于是就找到了Keras,一直用到现在。不仅仅是笔者,也许Keras的作者们当时都没有想到,在今天Keras居然成为了tensorflow的黄金标准。

笔者觉得这并非偶然。tensorflow曾有过不少上层API框架,比如tensorlayer,比如tf.slim,比如TFLearn,但为啥最终选择了Keras?除了Keras的“历史悠久”之外,还在于Keras真正当之无愧为一个优雅的封装实现。这一年来,笔者时不时会去读Keras源码,每次读Keras源码都会被它的严密性、优雅性所震撼,它是一件当之无愧的艺术品,一个人性化的创造。

所以,这是人性化的胜利。

转载到请包括本文地址:https://kexue.fm/archives/7055

更详细的转载事宜请参考:《科学空间FAQ》

如果您还有什么疑惑或建议,欢迎在下方评论区继续讨论。

如果您觉得本文还不错,欢迎分享/打赏本文。打赏并非要从中获得收益,而是希望知道科学空间获得了多少读者的真心关注。当然,如果你无视它,也不会影响你的阅读。再次表示欢迎和感谢!

如果您需要引用本文,请参考:

苏剑林. (Nov. 06, 2019). 《Keras:Tensorflow的黄金标准 》[Blog post]. Retrieved from https://kexue.fm/archives/7055

@online{kexuefm-7055,
        title={Keras:Tensorflow的黄金标准},
        author={苏剑林},
        year={2019},
        month={Nov},
        url={\url{https://kexue.fm/archives/7055}},
}