TensorFlow 2.0 是对 1.x 版本做了一次大的瘦身，Eager Execution 默认开启，并且使用 Keras 作为默认高级 API，
这些改进大大降低的 TensorFlow 使用难度。
本文主要记录了一次曲折的使用 Keras+TensorFlow2.0 的 BatchNormalization 的踩坑经历，这个坑差点要把 TF2.0 的新特性都毁灭殆尽，如果你在学习 TF2.0 的官方教程，不妨一观。
问题的产生
从教程[1]https://www.tensorflow.org/alpha/tutorials/images/transfer_learning?hl=zh-cn（讲述如何 Transfer Learning）说起：
IMG_SHAPE = (IMG_SIZE, IMG_SIZE, 3)
# Create the base model from the pre-trained model MobileNet V2
base_model = tf.keras.applications.MobileNetV2(input_shape=IMG_SHAPE,
include_top=False,weights='imagenet')
model = tf.keras.Sequential([
base_model,
tf.keras.layers.GlobalAveragePooling2D(),
tf.keras.layers.Dense(NUM_CLASSES)
])简单的代码我们就复用了 MobileNetV2 的结构创建了一个分类器模型，接着我们就可以调用 Keras 的接口去训练模型：
model.compile(optimizer=tf.keras.optimizers.RMSprop(lr=base_learning_rate),
loss='sparse_categorical_crossentropy',
metrics=['sparse_categorical_accuracy'])
model.summary()
history = model.fit(train_batches.repeat(),
epochs=20,
steps_per_epoch = steps_per_epoch,
validation_data=validation_batches.repeat(),
validation_steps=validation_steps)输出的结果看，一起都很完美：
Model: "sequential"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
mobilenetv2_1.00_160 (Model) (None, 5, 5, 1280) 2257984
_________________________________________________________________
global_average_pooling2d (Gl (None, 1280) 0
_________________________________________________________________
dense (Dense) (None, 2) 1281
=================================================================
Total params: 2,259,265
Trainable params: 1,281
Non-trainable params: 2,257,984
_________________________________________________________________
Epoch 11/20
581/581 [==============================] - 134s 231ms/step - loss: 0.4208 - accuracy: 0.9484 - val_loss: 0.1907 - val_accuracy: 0.9812
Epoch 12/20
581/581 [==============================] - 114s 197ms/step - loss: 0.3359 - accuracy: 0.9570 - val_loss: 0.1835 - val_accuracy: 0.9844
Epoch 13/20
581/581 [==============================] - 116s 200ms/step - loss: 0.2930 - accuracy: 0.9650 - val_loss: 0.1505 - val_accuracy: 0.9844
Epoch 14/20
581/581 [==============================] - 114s 196ms/step - loss: 0.2561 - accuracy: 0.9701 - val_loss: 0.1575 - val_accuracy: 0.9859
Epoch 15/20
581/581 [==============================] - 119s 206ms/step - loss: 0.2302 - accuracy: 0.9715 - val_loss: 0.1600 - val_accuracy: 0.9812
Epoch 16/20
581/581 [==============================] - 115s 197ms/step - loss: 0.2134 - accuracy: 0.9747 - val_loss: 0.1407 - val_accuracy: 0.9828
Epoch 17/20
581/581 [==============================] - 115s 197ms/step - loss: 0.1546 - accuracy: 0.9813 - val_loss: 0.0944 - val_accuracy: 0.9828
Epoch 18/20
581/581 [==============================] - 116s 200ms/step - loss: 0.1636 - accuracy: 0.9794 - val_loss: 0.0947 - val_accuracy: 0.9844
Epoch 19/20
581/581 [==============================] - 115s 198ms/step - loss: 0.1356 - accuracy: 0.9823 - val_loss: 0.1169 - val_accuracy: 0.9828
Epoch 20/20
581/581 [==============================] - 116s 199ms/step - loss: 0.1243 - accuracy: 0.9849 - val_loss: 0.1121 - val_accuracy: 0.9875然而这种写法还是不方便 Debug，我们希望可以精细的控制迭代的过程，并能够看到中间结果，所以我们训练的过程改成了这样：
optimizer = tf.keras.optimizers.RMSprop(lr=base_learning_rate)
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')
@tf.function
def train_cls_step(image, label):
with tf.GradientTape() as tape:
predictions = model(image)
loss = tf.keras.losses.SparseCategoricalCrossentropy()(label, predictions)
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
train_accuracy(label, predictions)
for images, labels in train_batches:
train_cls_step(images,labels)重新训练后，结果依然很完美！
但是，这时候我们想对比一下 Finetune 和重头开始训练的差别，所以把构建模型的代码改成了这样：
base_model = tf.keras.applications.MobileNetV2(input_shape=IMG_SHAPE,
include_top=False,weights=None)使得模型的权重随机生成，这时候训练结果就开始抽风了，Loss 不下降，Accuracy 稳定在 50% 附近游荡：
Step #10: loss=0.6937199831008911 acc=46.5625%
Step #20: loss=0.6932525634765625 acc=47.8125%
Step #30: loss=0.699873685836792 acc=49.16666793823242%
Step #40: loss=0.6910845041275024 acc=49.6875%
Step #50: loss=0.6935917139053345 acc=50.0625%
Step #60: loss=0.6965731382369995 acc=49.6875%
Step #70: loss=0.6949992179870605 acc=49.19642639160156%
Step #80: loss=0.6942993402481079 acc=49.84375%
Step #90: loss=0.6933775544166565 acc=49.65277862548828%
Step #100: loss=0.6928421258926392 acc=49.5%
Step #110: loss=0.6883170008659363 acc=49.54545593261719%
Step #120: loss=0.695658802986145 acc=49.453125%
Step #130: loss=0.6875559091567993 acc=49.61538314819336%
Step #140: loss=0.6851695775985718 acc=49.86606979370117%
Step #150: loss=0.6978713274002075 acc=49.875%
Step #160: loss=0.7165156602859497 acc=50.0%
Step #170: loss=0.6945627331733704 acc=49.797794342041016%
Step #180: loss=0.6936900615692139 acc=49.9305534362793%
Step #190: loss=0.6938323974609375 acc=49.83552551269531%
Step #200: loss=0.7030564546585083 acc=49.828125%
Step #210: loss=0.6926192045211792 acc=49.76190185546875%
Step #220: loss=0.6932414770126343 acc=49.786930084228516%
Step #230: loss=0.6924526691436768 acc=49.82337188720703%
Step #240: loss=0.6882281303405762 acc=49.869789123535156%
Step #250: loss=0.6877702474594116 acc=49.86249923706055%
Step #260: loss=0.6933954954147339 acc=49.77163314819336%
Step #270: loss=0.6944763660430908 acc=49.75694274902344%
Step #280: loss=0.6945018768310547 acc=49.49776840209961%我们将 predictions 的结果打印出来，发现 batch 内每个输出都是一模一样的：
0 = tf.Tensor([0.51352817 0.48647183], shape=(2,), dtype=float32)
1 = tf.Tensor([0.51352817 0.48647183], shape=(2,), dtype=float32)
2 = tf.Tensor([0.51352817 0.48647183], shape=(2,), dtype=float32)
3 = tf.Tensor([0.51352817 0.48647183], shape=(2,), dtype=float32)
4 = tf.Tensor([0.51352817 0.48647183], shape=(2,), dtype=float32)
5 = tf.Tensor([0.51352817 0.48647183], shape=(2,), dtype=float32)
6 = tf.Tensor([0.51352817 0.48647183], shape=(2,), dtype=float32)
7 = tf.Tensor([0.51352817 0.48647183], shape=(2,), dtype=float32)
8 = tf.Tensor([0.51352817 0.48647183], shape=(2,), dtype=float32)
9 = tf.Tensor([0.51352817 0.48647183], shape=(2,), dtype=float32)只是修改了初始权重，为何会产生这样的结果？
问题排查
实验 1
是不是训练不够充分，或者 learning rate 设置的不合适？
经过几轮调整，发现无论训练多久，learning rate 变大变小，都无法改变这种结果
实验 2
既然是权重的问题，是不是权重随机初始化的有问题，把初始权重拿出来统计了一下，一切正常
实验 3
这种问题根据之前的经验，在导出 Inference 模型的时候 BatchNormalization 没有处理好会出现这种一个 batch 内所有结果都一样的问题。但是如何解释训练的时候为什么会出现这个问题？而且为什么 Finetue 不会出现问题呢？只是改了权重的初始值而已呀
按照这个方向去 Google 的一番，发现了 Keras 的 BatchNormalization 确实有很多 issue，其中一个问题是在保存模型的是 BatchNormalzation 的 moving mean 和 moving variance 不会被保存 [6]https://github.com/tensorflow/tensorflow/issues/16455，而另外一个 issue 提到问题就和我们问题有关系的了：
[2] https://github.com/tensorflow/tensorflow/issues/19643
[3] https://github.com/tensorflow/tensorflow/issues/23873
最后，这位作者找到了原因，并且总结在了这里：
[4] https://pgaleone.eu/tensorflow/keras/2019/01/19/keras-not-yet-interface-to-tensorflow/
根据这个提示，我们做了如下尝试：
实验 3.1
改用 model.fit 的写法进行训练，在最初的几个 epoch 里面，我们发现好的一点的是 training accuracy 已经开始缓慢提升了，但是 validation accuracy 存在原来的问题。而且通过 model.predict_on_batch()拿到中间结果，发现依然还是 batch 内输出都一样。
Epoch 1/20
581/581 [==============================] - 162s 279ms/step - loss: 0.6768 - sparse_categorical_accuracy: 0.6224 - val_loss: 0.6981 - val_sparse_categorical_accuracy: 0.4984
Epoch 2/20
581/581 [==============================] - 133s 228ms/step - loss: 0.4847 - sparse_categorical_accuracy: 0.7684 - val_loss: 0.6931 - val_sparse_categorical_accuracy: 0.5016
Epoch 3/20
581/581 [==============================] - 130s 223ms/step - loss: 0.3905 - sparse_categorical_accuracy: 0.8250 - val_loss: 0.6996 - val_sparse_categorical_accuracy: 0.4984
Epoch 4/20
581/581 [==============================] - 131s 225ms/step - loss: 0.3113 - sparse_categorical_accuracy: 0.8660 - val_loss: 0.6935 - val_sparse_categorical_accuracy: 0.5016但是，随着训练的深入，结果出现了逆转，开始变得正常了（tf.function 的写法是无论怎么训练都不会变化，幸好没有放弃治疗）（追加：其实这里还是有问题的，继续看后面，当时就觉得怪怪的，不应该收敛这么慢）
Epoch 18/20
581/581 [==============================] - 131s 226ms/step - loss: 0.0731 - sparse_categorical_accuracy: 0.9725 - val_loss: 1.4896 - val_sparse_categorical_accuracy: 0.8703
Epoch 19/20
581/581 [==============================] - 130s 225ms/step - loss: 0.0664 - sparse_categorical_accuracy: 0.9748 - val_loss: 0.6890 - val_sparse_categorical_accuracy: 0.9016
Epoch 20/20
581/581 [==============================] - 126s 217ms/step - loss: 0.0631 - sparse_categorical_accuracy: 0.9768 - val_loss: 1.0290 - val_sparse_categorical_accuracy: 0.9031通多 model.predict_on_batch()拿到的结果也和这个 Accuracy 也是一致的
实验 3.2
通过上一个实验，我们验证了确实如果只通过 Keras 的 API 去训练，是正常。更深层的原因是什么呢？是不是 BatchNomalization 没有 update moving mean 和 moving variance 导致的呢？答案是 Yes
我们分别在两中训练方法前后，打印 moving mean 和 moving variance 的值：
def get_bn_vars(collection):
moving_mean, moving_variance = None, None for var in collection:
name = var.name.lower()
if "variance" in name:
moving_variance = var
if "mean" in name:
moving_mean = var
if moving_mean is not None and moving_variance is not None:
return moving_mean, moving_variance
raise ValueError("Unable to find moving mean and variance")
mean, variance = get_bn_vars(model.variables)
print(mean)
print(variance)我们发现，确实如果使用 model.fit()进行训练，mean 和 variance 是在 update 的 (虽然更新的速率看着有些奇怪)， 但是对于 tf.function 那种写法这两个值就没有被 update
那这里我们也可以解释为什么 Finetune 不会出现问题了，因为 imagenet 训练的mean, variance 已经是一个比较好的值了，即使不更新也可以正常使用
实验 3.3
是不是改成 [4] 里面说的方法构建动态的 Input_Shape 的模型就 OK 了呢？
class MyModel(Model):
def __init__(self):
super(MyModel, self).__init__()
self.conv1 = Conv2D(32, 3, activation='relu')
self.batch_norm1=BatchNormalization()
self.flatten = Flatten()
self.d1 = Dense(128, activation='relu')
self.d2 = Dense(10, activation='softmax')
def call(self, x):
x = self.conv1(x)
x = self.batch_norm1(x)
x = self.flatten(x)
x = self.d1(x)
return self.d2(x)
model = MyModel()
#model.build((None,28,28,1))
model.summary()
@tf.functiondef train_step(image, label):
with tf.GradientTape() as tape:
predictions = model(image)
loss = loss_object(label, predictions)
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
train_loss(loss)
train_accuracy(label, predictions)模型如下：
Model: "my_model"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
conv2d (Conv2D) multiple 320
_________________________________________________________________
batch_normalization_v2 (Batc multiple 128
_________________________________________________________________
flatten (Flatten) multiple 0
_________________________________________________________________
dense (Dense) multiple 2769024
_________________________________________________________________
dense_1 (Dense) multiple 1290
=================================================================
Total params: 2,770,762
Trainable params: 2,770,698
Non-trainable params: 64从 Output Shape 看，构建模型没问题
跑了一遍 MINST，结果也很不错！
以防万一，我们同样测试了一下 mean 和 variance 是否被更新，然而结果出乎意料，并没有！
也就是说 [4] 里面说的方案在我们这里并不可行
实验 3.4
既然我们定位问题是在 BatchNormalization 这里，所以就想到 BatchNormalization 的 training 和 testing 时候行为是不一致的，在 testing 的时候 moving mean 和 variance 是不需要 update 的，那么会不会是 tf.function 的这种写法并不会自动更改这个状态呢？
查看源码，发现 BatchNormalization 的 call()存在一个 training 参数，而且默认是 False
Call arguments:
inputs: Input tensor (of any rank).
training: Python boolean indicating whether the layer should behave in
training mode or in inference mode.
- `training=True`: The layer will normalize its inputs using the
mean and variance of the current batch of inputs.
- `training=False`: The layer will normalize its inputs using the
mean and variance of its moving statistics, learned during training.所以，做了如下改进：
class MyModel(Model):
def __init__(self):
super(MyModel, self).__init__()
self.conv1 = Conv2D(32, 3, activation='relu')
self.batch_norm1=BatchNormalization()
self.flatten = Flatten()
self.d1 = Dense(128, activation='relu')
self.d2 = Dense(10, activation='softmax')
def call(self, x,training=True):
x = self.conv1(x)
x = self.batch_norm1(x,training=training)
x = self.flatten(x)
x = self.d1(x)
return self.d2(x)
model = MyModel()
#model.build((None,28,28,1))
model.summary()
@tf.functiondef train_step(image, label):
with tf.GradientTape() as tape:
predictions = model(image,training=True)
loss = loss_object(label, predictions)
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
train_loss(loss)
train_accuracy(label, predictions)
@tf.functiondef test_step(image, label):
predictions = model(image,training=False)
t_loss = loss_object(label, predictions)
test_loss(t_loss)
test_accuracy(label, predictions)结果显示，moving mean 和 variance 开始更新啦，测试 Accuracy 也是符合预期
所以，我们可以确定问题的根源在于需要指定 BatchNormalization 是在 training 还是在 testing！
实验 3.5
3.4 中方法虽然解决了我们的问题，但是它是使用构建 Model 的 subclass 的方式，而我们之前的 MobileNetV2 是基于更加灵活 Keras Functional API 构建的，由于无法控制 call()函数的定义，没有办法灵活切换 training 和 testing 的状态，另外用 Sequential 的方式构建时也是一样。
[5]https://blog.keras.io/keras-as-a-simplified-interface-to-tensorflow-tutorial.html
[7]https://github.com/keras-team/keras/issues/7085
[8]https://github.com/keras-team/keras/issues/6752
从 5[8]中，我了解到两个情况，
-
- tf.keras.backend.set_learning_phase()可以改变 training 和 testing 的状态；
-
- model.updates 和 layer.updates 存着 old_value 和 new_value 的 Assign Op
所以我首先尝试：
tf.keras.backend.set_learning_phase(True)结果，MobileNetV2 构建的模型也可以正常工作了。
而且收敛的速度似乎比 model.fit()还快了很多，结合之前 model.fit()收敛慢的困惑，这里又增加的一个实验，在 model.fit()的版本里面也加上这句话，发现同样收敛速度也变快了！1 个 epoch 就能得到不错的结果了！
因此，这里又产生了一个问题 model.fit()到底有没有设 learning_phase 状态？如果没有是怎么做 moving mean 和 variance 的 update 的？
第二个方法，由于教程中讲述的是如何在 1.x 的版本构建，而在 eager execution 模式下，似乎没有办法去 run 这些 Assign Operation。仅做参考吧
update_ops = []
for assign_op in model.updates:
update_ops.append(assign_op))
#但是不知道拿到这些 update_ops 在 eager execution 模式下怎么处理呢？结论
总结一下，我们从 [4] 找到了解决问题的启发点，但是最终证明 [4] 里面的问题和解决方法用到我们这里并不能真正解决问题，问题的关键还是在于 Keras+TensorFlow2.0 里面我们如何处理在 training 和 testing 状态下行为不一致的 Layer；以及对于 model.fit()和 tf.funtion 这两种训练方法的区别，最终来看 model.fit()里面似乎包含很多诡异的行为。
最终的使用建议如下：
- 在使用 model.fit()或者 model.train_on_batch()这种 Keras 的 API 训练模型时，也推荐手动设置 tf.keras.backend.set_learning_phase(True)，可以加快收敛
- 如果使用 eager execution 这种方法，
- 1）使用构建 Model 的 subclass，但是针对 call()设置 training 的状态，对于 BatchNoramlization，Dropout 这样的 Layer 进行不同处理
- 2）使用 Functional API 或者 Sequential 的方式构建 Model，设置 tf.keras.backend.set_learning_phase(True)，但是注意在 testing 的时候改变一下状态
最后，为什么 TF 2.0 的教程里面没有提及这些？默认你已经精通 Keras 了吗？[捂脸哭]
感谢
感谢柏涛 帆月 应知老师提供的帮助
[1]https://www.tensorflow.org/alpha/tutorials/images/transfer_learning?hl=zh-cn
[2] https://github.com/tensorflow/tensorflow/issues/19643
[3] https://github.com/tensorflow/tensorflow/issues/23873
[4] https://pgaleone.eu/tensorflow/keras/2019/01/19/keras-not-yet-interface-to-tensorflow/
[5]https://blog.keras.io/keras-as-a-simplified-interface-to-tensorflow-tutorial.html
[6]https://github.com/tensorflow/tensorflow/issues/16455
[7]https://github.com/keras-team/keras/issues/7085
[8]https://github.com/keras-team/keras/issues/6752
本文作者：烁凡
阅读原文
本文为云栖社区原创内容，未经允许不得转载。