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

Question

我想在层x上执行一个简单的线性变换，因此变换的输出是y=a*x b。我正在处理图像，所以x是三维的（高度*宽度*通道）。然后，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。这些应该是模型在训练期间学习到的参数，但我不确定如何构造或指定它们。通常，我只是指定层（卷积层或密集层）来学习权重，但我不确定这里使用哪一层以及该层应该接受什么作为输入。我是否必须以某种方式初始化1的向量，然后学习权重以将其转换为gamma和beta？

即使可以使用TensorFlow的batchnorm层来实现这一点（了解它仍然很有用），我还是想了解如何从头开始实现这种缩放/转换。非常感谢。

Answer 1

正如注释中所提到的，可以通过自定义Keras层来实现这一点（请参阅有关这方面的教程）。可以对层基类进行子类化，并实现转换的行为。在下面的示例中，该层包含可学习的权重gamma和beta，它们都具有形状（num\u 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)