如何在张量流中指定线性变换？

Question

我想对层 x 执行一个简单的线性变换，所以变换的输出是y = a*x + b。我正在处理图像，所以 x 是 3 维的 (height * width * channels)。那么a 是大小为c 的尺度向量，其中c 是通道数，它对于x 的每个通道维度都有一个尺度参数。类似地，b 是一个大小的移位向量，它对于x 的每个通道维度都有一个移位参数。这是一个简单的归一化变体，没有对批次统计数据进行归一化。

这是一个例子：

# TODO: learn gamma and beta parameters
x = tf.keras.layers.Conv2D(filters=num_filters_e1,
    kernel_size=5,
    strides=2,
    padding='same')(input)
x = tf.keras.layers.Multiply()([x, gamma]) # scale by gamma along channel dim
x = tf.keras.layers.Add()([x, beta]) # shift with beta along channel dim
y = tf.keras.layers.ReLU()(x) # apply activation after transformation

我不确定如何获得 gamma 和 beta。这些应该是模型在训练期间学习的参数，但我不确定如何构造或指定它们。通常我只是指定层（卷积层或密集层）来学习权重，但我不确定在这里使用哪个层以及该层应该将什么作为输入。我是否必须以某种方式初始化一个向量，然后学习权重以将它们转换为 gamma 和 beta？

即使可以使用 TensorFlow 的 batchnorm layer 来做到这一点（这仍然很有用），我想学习如何从头开始实现这种缩放/移位。谢谢！

Answer 1

如 cmets 中所述，可以使用自定义 Keras 层 (see the tf.keras tutorial on this) 来实现这一点。可以子类化层基类并实现变换的行为。在下面的示例中，该层包含可学习的权重gamma 和beta，两者的形状均为(num_channels,)。权重分别用 1 和 0 初始化。

import tensorflow as tf

class LinearTransform(tf.keras.layers.Layer):
    """Layer that implements y=m*x+b, where m and b are
    learnable parameters.
    """
    def __init__(
        self,
        gamma_initializer="ones",
        beta_initializer="zeros",
        dtype=None,
        **kwargs
    ):
        super().__init__(dtype=dtype, **kwargs)
        self.gamma_initializer = gamma_initializer
        self.beta_initializer = beta_initializer

    def build(self, input_shape):
        num_channels = int(input_shape[-1])
        self.gamma = self.add_weight(
            "gamma",
            shape=[num_channels],
            initializer=self.gamma_initializer,
            dtype=self.dtype,
        )
        self.beta = self.add_weight(
            "beta",
            shape=[num_channels],
            initializer=self.beta_initializer,
            dtype=self.dtype,
        )

    def call(self, inputs):
        return self.gamma * inputs + self.beta

以下是行为测试：

tf.random.set_seed(42)
inputs = tf.random.normal([1, 24, 24, 4], dtype="float32", seed=42)

layer = LinearTransform()
np.testing.assert_allclose(layer(inputs), inputs)

layer = LinearTransform(
    gamma_initializer=tf.keras.initializers.constant(4),
    beta_initializer=tf.keras.initializers.constant(1),
)
np.testing.assert_allclose(layer(inputs), inputs * 4 + 1)