Keras 2.2.x learning_phase 机制源码级解读

由于Keras 2.3.0开始适配tf2.0导致代码大规模重构,因此我们讨论和tf1.x适配的Keras2.1.x与2.2.x版本。
Learning_phase虽然看似简单,实则非常重要,标志着模型的运行状态(训练还是推理)。在pytorch中可以使用model.train()与model.eval()切换模型的状态,因为会关系到BN与dropout层的计算改变。而在Keras中,为了使模型api简洁,且适配多种backend,运用了learning_phase机制去解决这种问题。而这种机制意外带来了巨大的复杂性,逻辑变得混乱,从而间接引发了一些问题:
Keras中的BN层错误:https://zhuanlan.zhihu.com/p/56225304
固化Keras模型后输出参数变了一点:https://stackoverflow.com/questions/61619032/got-small-output-value-error-between-h5-model-and-pb-model

本文试图彻底理清learning_phase机制,从而对Keras有着更全面的认知,使得下次遇到相关问题可以轻松解决。

一、定义

首先可以发现,在tensorflow_backend.py中有它的定义:

def learning_phase():
    """Returns the learning phase flag.

    The learning phase flag is a bool tensor (0 = test, 1 = train)
    to be passed as input to any Keras function
    that uses a different behavior at train time and test time.

    # Returns
        Learning phase (scalar integer tensor or Python integer).
    """
    graph = tf.get_default_graph()
    if graph not in _GRAPH_LEARNING_PHASES:
        phase = tf.placeholder_with_default(False,
                                            shape=(),
                                            name='keras_learning_phase')
        _GRAPH_LEARNING_PHASES[graph] = phase
    return _GRAPH_LEARNING_PHASES[graph]

可以看到,K.learning_phase()是全局量,且依附于当前graph中的唯一量,当调用它的时候,会去找_GRAPH_LEARNING_PHASES字典,若有当前graph的K.learning_phase()则取出,否则新建一个中placeholder_with_default(False)存入字典,下次取出的就是它了。
注意这是一个dtype=bool的placeholder,代表运行网络时,我们可以使用feed_dict={K.learning_phase(): 0 or 1}喂入,指定它的取值(前提是当前graph中的K.learning_phase()还是一个placeholder)。若不指定,则默认为False(0)。

  • 手动赋值
def set_learning_phase(value):
    """Sets the learning phase to a fixed value.

    # Arguments
        value: Learning phase value, either 0 or 1 (integers).

    # Raises
        ValueError: if `value` is neither `0` nor `1`.
    """
    global _GRAPH_LEARNING_PHASES
    if value not in {0, 1}:
        raise ValueError('Expected learning phase to be '
                         '0 or 1.')
    _GRAPH_LEARNING_PHASES[tf.get_default_graph()] = value

可以通过K.learning_phase()方法,为该全局量赋值,0为test,1为train。
赋值必须为这两个int之一。一旦手动赋值,则会覆盖掉原先dict里默认的placeholder。因此手动set值之后建立模型后,graph里的该值就被定死为一个int了,无法更改,也无法再feed_dict里喂入K.learning_phase()了。因为当前graph的模型建立时,用到的K.learning_phase()就是一个int。
建立模型后,再使用set_learning_phase更改值,对原模型无效。因为此时改的K.learning_phase()和graph里的建立好的模型无关了。

因此训练时候不能用,否则训练过程中预测验证集的时候也会强制使用train时候的配置。这与pytorch动态图能随意切换train/eval状态不同(毕竟tf是静态图)。

专门进行测试的时候可以用,但是其实没必要,一是因为Keras.model.predict的时候会传入数值0,二是因为placeholder_with_default默认就是False。

模型中不存在BN或dropout层的时候无效,因为train/eval都是同一套运算流程和参数配置。

那这个哪里可用到?可在建立模型的时候用,手动控制该layer使用哪种配置,适合折腾(比如开篇知乎那个解决BN问题的时候,就可通过创建layer时给某些BN层强制配置值0使得BN一直处于推理阶段,一直使用一开始迁移学习初始状态的移动平均值)

二、使用地点

在建立模型时,当遇到dropout或bn层时,以简单的keras.layers.Dropout()为例:

    def call(self, inputs, training=None):
        if 0. < self.rate < 1.:
            noise_shape = self._get_noise_shape(inputs)

            def dropped_inputs():
                return K.dropout(inputs, self.rate, noise_shape,
                                 seed=self.seed)
            return K.in_train_phase(dropped_inputs, inputs,
                                    training=training)
        return inputs

call()方法是所有Layer的逻辑实现层,调用该层的时候就会调用此方法。 K.dropout本质是为了适配不同后端,tf就会在该方法中调用tf.nn.dropout。

重点是 K.in_train_phase(dropped_inputs, inputs,training=training)函数。

def in_train_phase(x, alt, training=None):
    """Selects `x` in train phase, and `alt` otherwise.

    Note that `alt` should have the *same shape* as `x`.

    # Arguments
        x: What to return in train phase
            (tensor or callable that returns a tensor).
        alt: What to return otherwise
            (tensor or callable that returns a tensor).
        training: Optional scalar tensor
            (or Python boolean, or Python integer)
            specifying the learning phase.

    # Returns
        Either `x` or `alt` based on the `training` flag.
        the `training` flag defaults to `K.learning_phase()`.
    """
    if training is None:
        training = learning_phase()
        uses_learning_phase = True
    else:
        uses_learning_phase = False

    if training is 1 or training is True:
        if callable(x):
            return x()
        else:
            return x

    elif training is 0 or training is False:
        if callable(alt):
            return alt()
        else:
            return alt

    # else: assume learning phase is a placeholder tensor.
    x = switch(training, x, alt)
    if uses_learning_phase:
        x._uses_learning_phase = True
    return x

即通过K.in_train_phase判断该返回哪个值,训练的时候应该返回drop后的值,测试的时候应该不丢弃值(tf的dropout在训练时已经除以keep_prob系数了,所以测试时直接输出input即可)。

  1. 模型建立阶段:(即构建graph)layer未设置training参数,则默认为None。之前没手动set_learning_phase:
    则K.learning_phase()在training = learning_phase()这一步第一次被调用,然后初始化,故K.learning_phase()是:(BN的example)
    Tensor("bn_conv1/keras_learning_phase:0", shape=(), dtype=bool)
    当然若在模型建立前调用过,则name里没bn_conv1这个前缀。
    然后uses_learning_phase = True,代表这层使用到了learning_phase。返回的tensor一定会设置_uses_learning_phase = True这个属性值。graph建立完毕。
  • 若是训练阶段运行graph,Keras.model的fit等方法会自动判断模型是否有的layer._uses_learning_phase为True,即是否用到了learning_phase(),若用到了则feed_dict中多一个K.learning_phase()这个placeholder,并传入一个值:1. ,因此训练运行时,通过switch函数返回dropout的返回值。
  • 若测试阶段运行graph,Keras.model.predict()会自动为K.learning_phase()这个placeholder传入一个值:0.,或自行调用sess.run方法的时候不传值,由于placeholder_with_default(False)默认就是False值故不影响结果。

由于上文提到过,训练时一般不手动调用set_learning_phase,因此我们讨论:

  1. 模型建立阶段:(即构建graph)layer设置training参数为True,无论有没有手动set_learning_phase:
    我们发现graph搭建的时候直接跳过了swith函数分支,直接返回dropout之后的值。同时x._uses_learning_phase也未设置,可以说是完全抛弃了K.learning_phase()。K.learning_phase()都没初始化的机会。
  • 若是训练/测试阶段运行graph,都会直接运行graph里的返回dropout值,问题很大!
  1. 模型建立阶段:(即构建graph)layer设置training参数为False,无论之前有没有手动set_learning_phase:
    和2原理一样,都会强行返回输入值,graph中失去了对训练or测试阶段的选择性。即dropout层等于失效。

总结:
training参数手动设置True还是False都会带来巨变,使网络抛弃了K.learning_phase(),也等于是graph中失去了对训练or测试阶段的选择性。因为Graph可以应对两种阶段,本质是由于K.learning_phase()是个placeholder使输入有两种可能,然后K.in_train_phase中存在switch方法分支根据placeholder输出两种可能性。

那training参数保留None,手动set_learning_phase会咋样?

  1. 模型建立阶段:(即构建graph)layer未设置training参数,则默认为None。之前手动set_learning_phase=1:
    此时training=1,还是会进入分支2,只保留强行使用dropout这一条路。

  2. 模型建立阶段:(即构建graph)layer未设置training参数,则默认为None。之前手动set_learning_phase=0:
    此时training=0,还是会进入分支3,直接砍了dropout这一条路。即模型中dropout层没作用了。

总结:
构建模型时不应该手动设置training参数。那么,手动set_learning_phase=1构建模型,会使模型只留下训练配置一条路,和training=True一样,测试必定会强行使用dropout层一定输出会出错。而手动set_learning_phase=0构建模型,会使模型中dropout层失效,直接训练就废了。
那么还要手动设这些参数干嘛?因为有特殊情况,即迁移训练、或加载别人训练好的模型。例如keras.application.ResNet50(),我们不需要自己训练,那么可以提前set_learning_phase=0,这样构建出来的resnet会少了很多节点,相当于把graph里BN层训练的路子给砍了。load_weight的时候只会载入部分weights。model.predict()的时候结果也和原来一致。
同时,和知乎里说的一样,迁移训练时给BN设置training=False,或临时给BN层set_learning_phase=0,(别的层虽然set_learning_phase=1但由于他们训练or测试时行为一致所以其实无所谓),然后load_weights之后再训练,这样BN层只会输出旧模型的滑动平均值作为参数,都不会参与训练了。

三、训练或推理时

以keras.model.Model().predict()为例:

        # Prepare inputs, delegate logic to `_predict_loop`.
        if self.uses_learning_phase and not isinstance(K.learning_phase(), int):
            ins = x + [0.]
        else:
            ins = x
        self._make_predict_function()
        f = self.predict_function
        return self._predict_loop(f, ins, batch_size=batch_size,
                                  verbose=verbose, steps=steps)

而其中的self.uses_learning_phase来自model类继承的container类:

    @property
    def uses_learning_phase(self):
        return any([x._uses_learning_phase for x in self.outputs])

而这个方法是去判断container里的layer是否有含有x._uses_learning_phase属性。(BN、dropout层就有这个属性,上面提到了)
因此if self.uses_learning_phase and not isinstance(K.learning_phase(), int):这句的含义是若model(里的某些layer)用到了learning_phase,且当前graph的learning_phase()不是int,后面半句判断实际上是在判断是否model在创建之前使用过set_learning_phase,因为一旦set过0或1,那么模型实际上就会被剪掉train或test的分支,指定不同的learning_phase也就没了意义,实际上此时根本不能指定learning_phase()了,因为此时模型内部的K.learning_phase()就是一个int。
因此只有满足这两个条件(模型里用到了dropout或BN层、且模型创建前未set_learning_phase)才会给输入多添加一个值:0. 这个值代表的是feed_dict中喂给K.learning_phase()这个bool placeholder的值为0,代表测试阶段。然后graph运行的时候会进入推理分支。
self._make_predict_function()函数会动态创建实际上的预测函数,根据需不需要传入learning_phase创建不同的需要喂入的feed_dict。

训练阶段同理。

因此我们可以同样使用如下tf函数,来获取输出的值,分别是train分支与推理分支,结果与model.predict相同:

model_ = ResNet50(include_top=False, pooling='avg', weights='imagenet')
print(K.learning_phase())  # Tensor("bn_conv1/keras_learning_phase:0", shape=(), dtype=bool)
sess = K.get_session()
preds = sess.run(net_model.get_output_at(0), feed_dict={net_model.get_input_at(0): x_input, 
              sess.graph.get_tensor_by_name('bn_conv1/keras_learning_phase:0':1)})
print('before constantize output:', np.array(preds).squeeze()[:10])

preds = sess.run(net_model.get_output_at(0), feed_dict={net_model.get_input_at(0): x_input, 
              sess.graph.get_tensor_by_name('bn_conv1/keras_learning_phase:0':0)})
print('before constantize output:', np.array(preds).squeeze()[:10])  # 与model.predict相同

当然推理时可省略learning_phase的传入,因为这个placeholder默认就是False。
因此我们可通过下面语句获取中间某些node的输出值来调试网络:

preds = sess.run(sess.graph.get_tensor_by_name('bn_conv1/batchnorm/add_1:0'),     
                    feed_dict={net_model.get_input_at(0): x_input})
print('before constantize bn_conv1/batchnorm/add_1:0:', np.array(preds).squeeze()[0,0,:10])

注意,K.learning_phase()的name不是固定的,而是看第一次在哪里调用它,name的前缀会不同。这个案例中,在模型创建前并未调用过它,因此它是在第一个BN里才用到,那里的name_scope下第一次初始化创建,因此全名里带prefix是bn_conv1/keras_learning_phase:0。若在model创建前就调用过K.learning_phase()则模型里存储的该tensor.name=keras_learning_phase:0,模型创建的时候直接就去调用它了。

四、总结

Keras默认情况下K.learning_phase()返回一个全局的placeholder_with_default(False),Keras使用这个输入量来控制模型到底是train/eval阶段,关系到dropout和BN层的状态。

  • Keras.model.fit()等方法默认会构造一个feed_dict:{K.learning_phase():1}喂入模型,而predict等方法同理会喂入0,这样就告诉了BN层或dropout层此时应该使用graph里的哪个分支。
  • tf的静态图特性决定了必须使用placeholder这种机制创建模型后,才能根据输入量切换train/eval阶段。一旦提前set_learning_phase(1)将使创建出来的模型永远只拥有train这一个分支,后续将无法更改模型的分支,因为创建的模型里只有那一个分支。有分支的前提是K.learning_phase()是一个待输入量placeholder。
  • 不仅如此,将BN或dropout层的training参数设为True或False,同样也会引发这种现象,即设为True后构建的graph就只存在使用dropout层这一条路子,设为False则表示完全不用,将会忽略K.learning_phase()的取值。而BN的情况更加复杂,可参考文章开头的知乎链接。

你可能感兴趣的:(Keras 2.2.x learning_phase 机制源码级解读)