Keras/TensorFlow自定义权重约束：矩阵对角元为对应行最大值实现方法

实现逻辑 #

针对形状为2d × d的权重矩阵W，我们直接对前d行、后d行两个子矩阵的对角位置赋值为所在行的最大值即可，全程使用元素级运算，避免冗余的对角矩阵生成/加减操作，GPU并行效率极高。

代码实现 #

import tensorflow as tf
from tensorflow.keras.constraints import Constraint

class DiagRowMaxConstraint(Constraint):
    def __call__(self, w):
        d = tf.shape(w)[1]
        # 计算所有行的最大值，形状为 [2d, 1] 适配广播规则
        row_max = tf.reduce_max(w, axis=-1, keepdims=True)
        # 生成对角掩码：前d行[i,i]、后d行[d+i,i]位置为True，其余为False
        diag_mask = tf.eye(d, dtype=tf.bool)
        diag_mask = tf.concat([diag_mask, diag_mask], axis=0)
        # 对角位置替换为行最大值，其余位置保持原值
        return tf.where(diag_mask, row_max, w)

若你实际使用的权重矩阵存储形状为d × 2d（行对应特征维度、列对应变换维度），仅需调整row_max的计算轴为axis=0，diag_mask的拼接轴改为axis=1即可。

方案优势 #

对比你此前的实现，本方案计算开销更低：

• 仅需要1次全局行最大值计算、1次掩码拼接、1次条件赋值，所有操作均为可并行的元素级操作
• 避免了tf.linalg.diag、tf.linalg.diag_part这类需要额外生成对角矩阵的操作，减少显存占用和计算冗余
• 无需对权重矩阵做多次加减运算，进一步降低计算量

使用方法 #

在定义层时直接通过kernel_constraint参数传入即可，示例：

# 示例：d=32，权重形状为64×32的层使用该约束
d = 32
layer = tf.keras.layers.Dense(
    units=2*d,
    kernel_constraint=DiagRowMaxConstraint()
)

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