让Keras更酷一些：中间变量、权重滑动和安全生成器

作者丨苏剑林

单位丨追一科技

研究方向丨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,),