Пакетная нормализация с 3D-свертками в TensorFlow

Я реализую модель, основанную на 3D-свертках (для задачи, аналогичной распознаванию действий), и я хочу использовать нормализацию партии (см. [Иоффе и Szegedy 2015]). Я не мог найти никакого учебника, посвященного 3D-конверам, поэтому я делаю здесь короткий, который я хотел бы рассмотреть с вами.

Ниже приведен код TensorFlow r0.12, и он явно указывает переменные - я имею в виду, что я не использую tf.contrib.learn, кроме функции tf.contrib.layers.batch_norm(). Я делаю это как для лучшего понимания того, как все работает под капотом, так и для большей свободы реализации (например, сводки переменных).

Я вернусь к случаю трехмерной свертки, сначала написав пример для полностью связного слоя, затем для двумерной свертки и, наконец, для 3D-случая. Просматривая код, было бы замечательно, если бы вы могли проверить, все ли сделано правильно - код работает, но я не уверен на 100% о том, как я применяю нормализацию партии. Я заканчиваю этот пост более детальным вопросом.

import tensorflow as tf

# This flag is used to allow/prevent batch normalization params updates
# depending on whether the model is being trained or used for prediction.
training = tf.placeholder_with_default(True, shape=())

Корпус с полной связью (FC)

# Input.
INPUT_SIZE = 512
u = tf.placeholder(tf.float32, shape=(None, INPUT_SIZE))

# FC params: weights only, no bias as per [Ioffe & Szegedy 2015].
FC_OUTPUT_LAYER_SIZE = 1024
w = tf.Variable(tf.truncated_normal(
    [INPUT_SIZE, FC_OUTPUT_LAYER_SIZE], dtype=tf.float32, stddev=1e-1))

# Layer output with no activation function (yet).
fc = tf.matmul(u, w)

# Batch normalization.
fc_bn = tf.contrib.layers.batch_norm(
    fc,
    center=True,
    scale=True,
    is_training=training,
    scope='fc-batch_norm')

# Activation function.
fc_bn_relu = tf.nn.relu(fc_bn)
print(fc_bn_relu)  # Tensor("Relu:0", shape=(?, 1024), dtype=float32)

Двумерный сверточный (CNN) слой

# Input: 640x480 RGB images (whitened input, hence tf.float32).
INPUT_HEIGHT = 480
INPUT_WIDTH = 640
INPUT_CHANNELS = 3
u = tf.placeholder(tf.float32, shape=(None, INPUT_HEIGHT, INPUT_WIDTH, INPUT_CHANNELS))

# CNN params: wights only, no bias as per [Ioffe & Szegedy 2015].
CNN_FILTER_HEIGHT = 3  # Space dimension.
CNN_FILTER_WIDTH = 3  # Space dimension.
CNN_FILTERS = 128
w = tf.Variable(tf.truncated_normal(
    [CNN_FILTER_HEIGHT, CNN_FILTER_WIDTH, INPUT_CHANNELS, CNN_FILTERS],
    dtype=tf.float32, stddev=1e-1))

# Layer output with no activation function (yet).
CNN_LAYER_STRIDE_VERTICAL = 1
CNN_LAYER_STRIDE_HORIZONTAL = 1
CNN_LAYER_PADDING = 'SAME'
cnn = tf.nn.conv2d(
    input=u, filter=w,
    strides=[1, CNN_LAYER_STRIDE_VERTICAL, CNN_LAYER_STRIDE_HORIZONTAL, 1],
    padding=CNN_LAYER_PADDING)

# Batch normalization.
cnn_bn = tf.contrib.layers.batch_norm(
    cnn,
    data_format='NHWC',  # Matching the "cnn" tensor which has shape (?, 480, 640, 128).
    center=True,
    scale=True,
    is_training=training,
    scope='cnn-batch_norm')

# Activation function.
cnn_bn_relu = tf.nn.relu(cnn_bn)
print(cnn_bn_relu)  # Tensor("Relu_1:0", shape=(?, 480, 640, 128), dtype=float32)

Случай трехмерного сверточного (CNN3D) слоя

# Input: sequence of 9 160x120 RGB images (whitened input, hence tf.float32).
INPUT_SEQ_LENGTH = 9
INPUT_HEIGHT = 120
INPUT_WIDTH = 160
INPUT_CHANNELS = 3
u = tf.placeholder(tf.float32, shape=(None, INPUT_SEQ_LENGTH, INPUT_HEIGHT, INPUT_WIDTH, INPUT_CHANNELS))

# CNN params: wights only, no bias as per [Ioffe & Szegedy 2015].
CNN3D_FILTER_LENGHT = 3  # Time dimension.
CNN3D_FILTER_HEIGHT = 3  # Space dimension.
CNN3D_FILTER_WIDTH = 3  # Space dimension.
CNN3D_FILTERS = 96
w = tf.Variable(tf.truncated_normal(
    [CNN3D_FILTER_LENGHT, CNN3D_FILTER_HEIGHT, CNN3D_FILTER_WIDTH, INPUT_CHANNELS, CNN3D_FILTERS],
    dtype=tf.float32, stddev=1e-1))

# Layer output with no activation function (yet).
CNN3D_LAYER_STRIDE_TEMPORAL = 1
CNN3D_LAYER_STRIDE_VERTICAL = 1
CNN3D_LAYER_STRIDE_HORIZONTAL = 1
CNN3D_LAYER_PADDING = 'SAME'
cnn3d = tf.nn.conv3d(
    input=u, filter=w,
    strides=[1, CNN3D_LAYER_STRIDE_TEMPORAL, CNN3D_LAYER_STRIDE_VERTICAL, CNN3D_LAYER_STRIDE_HORIZONTAL, 1],
    padding=CNN3D_LAYER_PADDING)

# Batch normalization.
cnn3d_bn = tf.contrib.layers.batch_norm(
    cnn3d,
    data_format='NHWC',  # Matching the "cnn" tensor which has shape (?, 9, 120, 160, 96).
    center=True,
    scale=True,
    is_training=training,
    scope='cnn3d-batch_norm')

# Activation function.
cnn3d_bn_relu = tf.nn.relu(cnn3d_bn)
print(cnn3d_bn_relu)  # Tensor("Relu_2:0", shape=(?, 9, 120, 160, 96), dtype=float32)

То, что я хотел бы убедиться, заключается в том, что код выше точно реализует нормализацию партии, как описано в [Ioffe and Szegedy 2015] в конце п. 3,2:

Для сверточных слоев мы дополнительно хотим, чтобы нормализация подчинялась сверточному свойству, так что различные элементы одной и той же карты признаков в разных местах нормализуются одинаково. Для этого мы совместно нормализуем все активации в мини-баре, во всех местах. [...] Алг. 2 модифицируется аналогично, так что во время вывода преобразование BN применяет одно и то же линейное преобразование к каждой активации в данной карте характеристик.

UPDATE Я думаю, что приведенный выше код также верен для случая 3D conv. Фактически, когда я определяю свою модель, если я печатаю все обучаемые переменные, я также вижу ожидаемые числа бета-и гамма-переменных. Например:

Tensor("conv3a/conv3d_weights/read:0", shape=(3, 3, 3, 128, 256), dtype=float32)
Tensor("BatchNorm_2/beta/read:0", shape=(256,), dtype=float32)
Tensor("BatchNorm_2/gamma/read:0", shape=(256,), dtype=float32)

Это выглядит хорошо для меня, поскольку из-за BN одна пара бета и гамма изучается для каждой карты характеристик (всего 256).

[Иоффе и Сегеды, 2015]: Пакетная нормализация: ускорение глубокой сетевой тренировки путем сокращения внутреннего сдвига ковариации

Корпус с полной связью (FC)

Двумерный сверточный (CNN) слой

Случай трехмерного сверточного (CNN3D) слоя

Ответ 1