如何在Keras中实现带权重相等约束的自定义层？

我明白你想要实现一个满足 y = Wx + c 的自定义层，其中W是分块方阵，c是元素重复的偏置向量。直接用权重约束来指定索引相等确实可行，但更高效且直观的方式是直接构造符合结构的权重和偏置，而不是依赖约束函数（约束会在每次梯度更新后修正，额外消耗计算资源）。下面我给你一步步拆解实现思路和代码：

核心思路 #

1. 分块矩阵W：不用直接初始化完整的W，而是初始化组成分块的子方阵（比如A、B），然后在前向传播时拼接成目标W矩阵。这样天然保证W的分块结构，没有多余参数。
2. 重复偏置c：初始化少量标量偏置（比如对应重复块的基值），然后通过重复扩展得到完整的偏置向量，确保元素重复的约束。

代码实现示例 #

假设我们的需求是：输入/输出维度为 2n，W的分块结构为：

[ A  B ]
[ B  A ]

其中A、B都是n×n的方阵；偏置c是前n个元素为c0、后n个元素为c1的重复向量。

import tensorflow as tf
from tensorflow.keras.layers import Layer

class CustomBlockLayer(Layer):
    def __init__(self, sub_matrix_dim, **kwargs):
        # sub_matrix_dim 是子方阵A、B的维度，输入/输出维度为 2*sub_matrix_dim
        self.sub_dim = sub_matrix_dim
        super().__init__(**kwargs)
    
    def build(self, input_shape):
        # 校验输入维度是否符合预期
        input_feature_dim = input_shape[-1]
        assert input_feature_dim == 2 * self.sub_dim, \
            f"输入特征维度必须为 {2*self.sub_dim}，当前为 {input_feature_dim}"
        
        # 初始化分块子方阵A和B
        self.A = self.add_weight(
            shape=(self.sub_dim, self.sub_dim),
            initializer='glorot_uniform',  # 常用的权重初始化器
            trainable=True,
            name='sub_matrix_A'
        )
        self.B = self.add_weight(
            shape=(self.sub_dim, self.sub_dim),
            initializer='glorot_uniform',
            trainable=True,
            name='sub_matrix_B'
        )
        
        # 初始化偏置的基标量（对应重复块的取值）
        self.c0 = self.add_weight(
            shape=(),  # 标量
            initializer='zeros',
            trainable=True,
            name='bias_c0'
        )
        self.c1 = self.add_weight(
            shape=(),
            initializer='zeros',
            trainable=True,
            name='bias_c1'
        )
        
        super().build(input_shape)  # 必须调用父类的build方法
    
    def call(self, inputs):
        # 1. 拼接构造分块矩阵W
        top_row = tf.concat([self.A, self.B], axis=1)
        bottom_row = tf.concat([self.B, self.A], axis=1)
        W = tf.concat([top_row, bottom_row], axis=0)
        
        # 2. 计算 Wx
        wx = tf.matmul(inputs, W)
        
        # 3. 构造重复偏置向量c
        c = tf.concat(
            [tf.repeat([self.c0], self.sub_dim), tf.repeat([self.c1], self.sub_dim)],
            axis=0
        )
        
        # 4. 输出 y = Wx + c
        return wx + c
    
    def compute_output_shape(self, input_shape):
        # 定义输出维度，可选（TensorFlow 2.x 通常能自动推断）
        return (input_shape[0], 2 * self.sub_dim)

灵活调整结构 #

如果你的分块结构或偏置重复规则不同，只需要修改call方法中的拼接逻辑：

• 比如W的结构是[[A, A], [B, B]]，只需把拼接W的代码改成：

  top_row = tf.concat([self.A, self.A], axis=1)
  bottom_row = tf.concat([self.B, self.B], axis=1)
  W = tf.concat([top_row, bottom_row], axis=0)
  

• 如果偏置是整个向量重复同一个标量，只需初始化一个c_scalar，然后：

  c = tf.repeat([self.c_scalar], 2 * self.sub_dim)
  

关于权重约束的备选方案 #

如果你确实想用约束函数来实现（比如某些复杂场景下无法直接构造），这里也给你一个示例：
比如强制W的分块对称，定义约束函数后在初始化W时指定：

def block_weight_constraint(W):
    n = W.shape[0] // 2
    # 取对称块的平均值来保证结构
    A = (W[:n, :n] + W[n:, n:]) / 2
    B = (W[:n, n:] + W[n:, :n]) / 2
    # 重构符合要求的W
    top_row = tf.concat([A, B], axis=1)
    bottom_row = tf.concat([B, A], axis=1)
    return tf.concat([top_row, bottom_row], axis=0)

# 在build中初始化W时指定约束
self.W = self.add_weight(
    shape=(2*self.sub_dim, 2*self.sub_dim),
    initializer='glorot_uniform',
    trainable=True,
    name='W',
    constraint=block_weight_constraint
)

但这种方式每次梯度更新后都会执行约束修正，效率不如直接构造子矩阵的方法，所以优先推荐前者。

内容的提问来源于stack exchange，提问作者Jonathan Lym