2021-08-18使用TensorFlow图层的“kernel_constraint”实现权重规范化

在[1]，Salimans等人 . 提出了一种神经网络归一化技术，称为权重归一化，它对网络层的权重向量进行归一化，与例如批量归一化相反，后者规范化流经该层的实际数据批量 . 在某些情况下，权重归一化方法的计算开销较低，并且还可以在使用批量归一化不可行的情况下使用 .

由优化器更新后应用于内核的可选投影函数（例如，用于实现层权重的范数约束或值约束） .

我的问题是：是否可以使用上述TensorFlow图层来实现权重标准化 kernel_constraint ？假设 x 是一个形状为 (batch, height, width, channels) 的输入，我想我可以按如下方式实现它：

x = tf.layers.conv2d(
    inputs=x,
    filters=16,
    kernel_size=(3, 3),
    strides=(1, 1),
    kernel_constraint=lambda kernel: (
        tf.nn.l2_normalize(w, list(range(kernel.shape.ndims-1)))))

tf.keras.constraints.UnitNorm(
    axis=0
)

# define model
model = Sequential()
model.add(Dense(500, input_dim=2, activation='relu', kernel_constraint=unit_norm()))