将 tf.contrib.layers.batch_norm 迁移到 TensorFlow 2.0

Question

我正在将 TensorFlow 代码迁移到 TensorFlow 2.1.0。

这是原始代码：

conv = tf.layers.conv2d(inputs, out_channels, kernel_size=3, padding='SAME')
conv = tf.contrib.layers.batch_norm(conv, updates_collections=None, decay=0.99, scale=True, center=True)
conv = tf.nn.relu(conv)
conv = tf.contrib.layers.max_pool2d(conv, 2)

这就是我所做的：

conv1 = Conv2D(out_channels, (3, 3), activation='relu', padding='same', data_format='channels_last', name=name)(inputs)
conv1 = Conv2D(64, (5, 5), activation='relu', padding='same', data_format="channels_last")(conv1)
#conv = tf.contrib.layers.batch_norm(conv, updates_collections=None, decay=0.99, scale=True, center=True)
pool1 = MaxPooling2D(pool_size=(2, 2), data_format="channels_last")(conv1)

我的问题是我不知道如何处理tf.contrib.layers.batch_norm。

如何将 tf.contrib.layers.batch_norm 迁移到 Tensorflow 2.x？

更新：
使用评论建议，我认为我已经正确迁移：

conv1 = BatchNormalization(momentum=0.99, scale=True, center=True)(conv1)

但我不确定decay 是否像momentum 并且我不知道如何在BatchNormalization 方法中设置updates_collections。

Answer 1

我在使用要微调的经过训练的模型时遇到了这个问题。只是像 OP 一样将 tf.contrib.layers.batch_norm 替换为 tf.keras.layers.BatchNormalization 确实给了我一个错误，其修复方法如下所述。

旧代码如下所示：

tf.contrib.layers.batch_norm(
    tensor,
    scale=True,
    center=True,
    is_training=self.use_batch_statistics,
    trainable=True,
    data_format=self._data_format,
    updates_collections=None,
)

更新后的工作代码如下所示：

tf.keras.layers.BatchNormalization(
    name="BatchNorm",
    scale=True,
    center=True,
    trainable=True,
)(tensor)

我不确定我删除的所有关键字参数是否都会出现问题，但似乎一切正常。注意name="BatchNorm" 参数。图层使用不同的命名架构，因此我不得不使用inspect_checkpoint.py 工具查看模型并找到恰好是BatchNorm 的图层名称。