python - tensorflow多GPU代码如何实现批量归一化层

Question

我创建了一个多 GPU 网络作为Cifar10_multigpu

在推理实现中，他们说：

我们使用 tf.get_variable() 而不是 tf.Variable() 实例化所有变量，以便在多个 GPU 训练运行中共享变量。如果我们只在单个 GPU 上运行这个模型，我们可以通过用 tf.Variable() 替换所有 tf​​.get_variable() 实例来简化这个函数。

所以我以我所有的 conv2d 层为例，但是 batchnorm 层呢？我如何自己实现它？

我可以tensorflow.contrib.slim.batch_norm在这种情况下使用吗？该示例不包含任何关于批规范层的建议。

Answer 1

只需使用tf.layers.batch_normalization. 它还通过 创建变量tf.get_variable()，因此它们也可以共享。

此外，它还可以与tf.layers.conv*功能无缝协作。

更新：tf.nn.batch_normalization也很好。这是一个更底层的功能，需要您自己管理mean和variance张量。实际上，tf.layers.batch_normalization是对函数的包装tf.nn.*，其中还包括tf.nn.fused_batch_norm（更快的融合版本）。