如何在顺序层中使用批量归一化层的训练选项？

Question

我正在尝试将批处理规范化放在顺序层中以获得干净的代码。

但是我发现，如果将其放入顺序中，我将不知道如何使用批处理规范化层的训练选项。

例如：

class OriginalNetwork(tf.keras.Model):
    def __init__(self):
        self.conv = tf.keras.layers.Conv2D(32,3,strides=2,padding='same')
        self.bm = tf.keras.layers.BatchNormalization()
        self.relu = tf.kears.layers.ReLU()
    def call(self, x, training=True):
        x = self.conv(x)
        x = self.bm(x, training=training)
        return self.relu(x)

我要使用以下结构代替OriginalNetwork。

def conv_block(channels, kernel_size, strides):
    layer = tf.keras.Sequential()
    layer.add(tf.keras.layers.Conv2D(channels, kernel_size, strides=strides, padding='same'))
    layer.add(tf.keras.layers.BatchNormalization())
    layer.add(tf.keras.layers.ReLU())
    return layer

class CleanNetwork(tf.keras.Model):
    def __init__(self):
        self.net = conv_block(32,3,2)
    def call(self, x, training=True):
        return self.net(x) # How to put training option??

但是，如何给批归一化层提供训练选项呢？