Пакетная нормализация с помощью трехмерных сверток в TensorFlow

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

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

Я плавно перейду к случаю трехмерной свертки, написав сначала пример для полносвязного слоя, затем для двумерной свертки и, наконец, для трехмерного случая. Просматривая код, было бы здорово, если бы вы могли проверить, все ли сделано правильно - код работает, но я не на 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 & Szegedy 2015] в конце разд. 3.2:

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

ОБНОВЛЕНИЕ. Я полагаю, что приведенный выше код также верен для случая трехмерной конвекции. Фактически, когда я определяю свою модель, если я печатаю все обучаемые переменные, я также вижу ожидаемое количество бета- и гамма-переменных. Например:

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).

[Ioffe & Szegedy 2015]: Пакетная нормализация: ускорение глубокого обучения сети за счет уменьшения внутреннего ковариантного сдвига