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

Question

I'm currently learning to use Tensorflow-2.0 in my project. I want to use a convolution neural network (CNN) to accomplish a semantic segmentation task and find a strange error when coding.

First of all, a simple model was constructed and work fine.

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)

However, When I try to use customize Keras Model, some error occurred:

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)

The error log is:

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

And I found that if I remove BatchNormalization layers in the call function, the code would work fine:

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)

So maybe the error is about the improper use of BatchNormalization layer. My TensorFlow version is 2.0.0-beta1. Why does this error happen? How can I fix this error? Thank you for your help!

Answer 1

tf.keras.layers.BatchNormalization is a trainable layer meaning it has parameters which will be updated during backward pass (namely gamma and beta corresponding to learned variance and mean for each feature).

In order for the gradient to be propagated, this layer has to be registered in Tensorflow's graph. This operation is done inside __init__ , when you assign to self , hence if you create this layers inside call it will not be registered correctly.

Code which should be working correctly:

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)

Additionally I have used Sequential so the operations are better kept together.