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

How to specify linear transformation in tensorflow?

我想对x层做一个简单的线性变换,所以变换的输出是y = a*x + b。我正在处理图像,所以 x 是 3 维的 (height * width * channels)。那么 a 是大小为 c 的尺度向量,其中 c 是通道数,并且它对于 x 的每个通道维度都有一个单一的尺度参数。类似地,b 是一个大小为 shift 的向量,它对于 x 的每个通道维度都有一个单独的 shift 参数。这是标准化的简单变体,没有标准化批处理统计信息。

这是一个例子:

# 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 可以做到这一点(了解它仍然很有用),我也想学习如何从头开始实现这个 scaling/shifting。谢谢!

如评论中所述,可以使用自定义 Keras 层 (see the tf.keras tutorial on this) 来完成此操作。可以子 class 层基础 class 并实现变换的行为。在下面的示例中,该层包含可学习的权重 gammabeta,它们的形状均为 (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)