python - 如何在自定义 Keras 模型中使用 BatchNormalization 层

Question

我目前正在学习在我的项目中使用 Tensorflow-2.0。我想使用卷积神经网络(CNN)完成一个语义分割任务，在编码时发现一个奇怪的错误。

首先，构建了一个简单的模型并且运行良好。

X_train,y_train = load_data()

input = tf.keras.layers.Input((512,512,7))
c1 = tf.keras.layers.Conv2D(64,3,padding='same',activation='relu')(input)
c1 = tf.keras.layers.BatchNormalization()(c1)
c1 = tf.keras.layers.Conv2D(64,3,padding='same',activation='relu')(c1)
c1 = tf.keras.layers.BatchNormalization()(c1)
c1 = tf.keras.layers.Conv2D(3,3,padding='same',activation='softmax')(c1)
model = tf.keras.models.Model(inputs=[input],outputs=[c1])

model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001),
              loss=tf.keras.losses.sparse_categorical_crossentropy,
              metrics=['accuracy'])
results = model.fit(X_train,y_train,batch_size=8,epochs=1000)

但是，当我尝试使用自定义 Keras 模型时，出现了一些错误:

class SequenceEECNN(tf.keras.Model):
    def __init__(self,n_class=3,width=32):
        super(SequenceEECNN,self).__init__(name='SequenceEECNN')
        self.n_class = n_class
        self.width = width
        self.c1 = tf.keras.layers.Conv2D(self.width, 3,activation='relu',padding='same')
        self.c2 = tf.keras.layers.Conv2D(self.width, 3, activation='relu',padding='same')
        self.out = tf.keras.layers.Conv2D(self.n_class,3,activation='softmax',padding='same')

    def call(self, inputs):
        x = self.c1(inputs)
        x = tf.keras.layers.BatchNormalization()(x)
        x = self.c2(x)
        x = tf.keras.layers.BatchNormalization()(x)
        return self.out(x)

X_train,y_train = load_data()

model = SequenceEECNN()

model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001),
              loss=tf.keras.losses.sparse_categorical_crossentropy,
              metrics=['accuracy'])
results = model.fit(X_train,y_train,batch_size=8,epochs=1000)

错误日志是:

Instructions for updating:
Use tf.where in 2.0, which has the same broadcast rule as np.where
Train on 128 samples
Epoch 1/1000
2019-08-11 16:21:27.377452: I tensorflow/stream_executor/platform/default/dso_loader.cc:42] Successfully opened dynamic library libcudnn.so.7
2019-08-11 16:21:27.378136: W tensorflow/core/framework/op_kernel.cc:1546] OP_REQUIRES failed at resource_variable_ops.cc:268 : Not found: Resource localhost/_AnonymousVar10/N10tensorflow3VarE does not exist.
2019-08-11 16:21:27.378156: W tensorflow/core/common_runtime/base_collective_executor.cc:216] BaseCollectiveExecutor::StartAbort Not found: Resource localhost/_AnonymousVar10/N10tensorflow3VarE does not exist.
     [[{{node Adam/gradients/SequenceEECNN/batch_normalization_1/cond_grad/If/then/_52/VariableShape_1}}]]
     [[Func/Adam/gradients/SequenceEECNN/batch_normalization/cond_grad/If/else/_75/input/_230/_72]]
2019-08-11 16:21:27.378314: W tensorflow/core/common_runtime/base_collective_executor.cc:216] BaseCollectiveExecutor::StartAbort Not found: Resource localhost/_AnonymousVar10/N10tensorflow3VarE does not exist.
     [[{{node Adam/gradients/SequenceEECNN/batch_normalization_1/cond_grad/If/then/_52/VariableShape_1}}]]
2019-08-11 16:21:27.378322: W tensorflow/core/framework/op_kernel.cc:1546] OP_REQUIRES failed at resource_variable_ops.cc:268 : Not found: Resource localhost/_AnonymousVar11/N10tensorflow3VarE does not exist.
Traceback (most recent call last):
  File "/media/xrzhang/Data/ZHS/Research/CNN-TF2/learn_tf2/test_model.py", line 40, in <module>
    results = model.fit(X_train,y_train,batch_size=8,epochs=1000)
  File "/media/xrzhang/Data/ZHS/Research/CNN-TF2/venv/lib/python3.6/site-packages/tensorflow/python/keras/engine/training.py", line 643, in fit
    use_multiprocessing=use_multiprocessing)
  File "/media/xrzhang/Data/ZHS/Research/CNN-TF2/venv/lib/python3.6/site-packages/tensorflow/python/keras/engine/training_arrays.py", line 664, in fit
    steps_name='steps_per_epoch')
  File "/media/xrzhang/Data/ZHS/Research/CNN-TF2/venv/lib/python3.6/site-packages/tensorflow/python/keras/engine/training_arrays.py", line 383, in model_iteration
    batch_outs = f(ins_batch)
  File "/media/xrzhang/Data/ZHS/Research/CNN-TF2/venv/lib/python3.6/site-packages/tensorflow/python/keras/backend.py", line 3510, in __call__
    outputs = self._graph_fn(*converted_inputs)
  File "/media/xrzhang/Data/ZHS/Research/CNN-TF2/venv/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 572, in __call__
    return self._call_flat(args)
  File "/media/xrzhang/Data/ZHS/Research/CNN-TF2/venv/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 671, in _call_flat
    outputs = self._inference_function.call(ctx, args)
  File "/media/xrzhang/Data/ZHS/Research/CNN-TF2/venv/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 445, in call
    ctx=ctx)
  File "/media/xrzhang/Data/ZHS/Research/CNN-TF2/venv/lib/python3.6/site-packages/tensorflow/python/eager/execute.py", line 67, in quick_execute
    six.raise_from(core._status_to_exception(e.code, message), None)
  File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.NotFoundError: 2 root error(s) found.
  (0) Not found:  Resource localhost/_AnonymousVar10/N10tensorflow3VarE does not exist.
     [[{{node Adam/gradients/SequenceEECNN/batch_normalization_1/cond_grad/If/then/_52/VariableShape_1}}]]
     [[Func/Adam/gradients/SequenceEECNN/batch_normalization/cond_grad/If/else/_75/input/_230/_72]]
  (1) Not found:  Resource localhost/_AnonymousVar10/N10tensorflow3VarE does not exist.
     [[{{node Adam/gradients/SequenceEECNN/batch_normalization_1/cond_grad/If/then/_52/VariableShape_1}}]]
0 successful operations.
0 derived errors ignored. [Op:__inference_keras_scratch_graph_1409]

Function call stack:
keras_scratch_graph -> keras_scratch_graph

而且我发现，如果我删除调用函数中的 BatchNormalization 层，代码就会正常工作:

class SequenceEECNN(tf.keras.Model):
    def __init__(self,n_class=3,width=32):
        super(SequenceEECNN,self).__init__(name='SequenceEECNN')
        self.n_class = n_class
        self.width = width
        self.c1 = tf.keras.layers.Conv2D(self.width, 3,activation='relu',padding='same')
        self.c2 = tf.keras.layers.Conv2D(self.width, 3, activation='relu',padding='same')
        self.out = tf.keras.layers.Conv2D(self.n_class,3,activation='softmax',padding='same')

    def call(self, inputs):
        x = self.c1(inputs)
        # x = tf.keras.layers.BatchNormalization()(x) remove any BatchNorm layer
        x = self.c2(x)
        x = tf.keras.layers.BatchNormalization()(x)
        return self.out(x)

所以这个错误可能是关于 BatchNormalization 层使用不当造成的。我的TensorFlow版本是2.0.0-beta1。为什么会出现这个错误？我该如何修复这个错误？感谢您的帮助!

Answer 1

tf.keras.layers.BatchNormalization 是一个可训练层，这意味着它具有将在向后传递期间更新的参数(即 gamma 和beta 对应于每个特征的学习方差和平均值)。

为了传播梯度，必须在 Tensorflow 图中注册该层。当您分配给 self 时，此操作是在 __init__ 内部完成的，因此，如果您在 call 内部创建此层，它将无法正确注册。

应该正常工作的代码:

class SequenceEECNN(tf.keras.Model):
    def __init__(self, n_class=3, width=32):
        super().__init__()
        self.n_class = n_class
        self.width = width
        self.first = tf.keras.Sequential(
            [
                tf.keras.layers.Conv2D(
                    self.width, 3, activation="relu", padding="same"
                ),
                tf.keras.layer.BatchNormalization(),
            ]
        )
        self.second = tf.keras.Sequential(
            [
                tf.keras.layers.Conv2D(
                    self.width, 3, activation="relu", padding="same"
                ),
                tf.keras.layer.BatchNormalization(),
            ]
        )
        self.out = tf.keras.layers.Conv2D(
            self.n_class, 3, activation="softmax", padding="same"
        )

    def call(self, inputs):
        x = self.first(inputs)
        x = self.second(x)
        return self.out(x)

此外，我还使用了Sequential，因此操作可以更好地保持在一起。