如何通过掩码冻结“Tensorflow”张量中的部分权重？

如何解决如何通过掩码冻结“Tensorflow”张量中的部分权重？

我想将某些部分的权重冻结为零，这样它们就不会在梯度下降过程中更新。当然，有一个选项可以自定义 fit 方法 对于 tf.keras.Model，并将梯度乘以掩码，但似乎应该有更短的解决方案。

在 Tensorflow 中有一个函数 stop_gradient，我尝试执行以下操作：

def freeze_weights_by_mask(tensor,mask):
    if mask is not None:
        z1 = tf.zeros_like(tensor)
        return tf.where(mask,tensor,z1) + tf.stop_gradient(tf.where(tf.math.logical_not(mask),z1,z1))
    return tensor  

我期望这个函数将一些权重设置为零，并冻结梯度。 然而，做完之后

model.weights[0] = mask_gradient(model.weights[0],mask)

并查看通过以下方式获得的 grad：

loss_fn = ...

with tf.GradientTape() as tape:
    logits = model(train_data[0][:32,:,:],training=True)

    loss_value = loss_fn(train_labels[0][:32,logits)

    grads = tape.gradient(loss_value,model.trainable_variables)

我发现 model.weights[0] 的所有梯度都不为零。

如何在短时间内完成？

在这个问题中，How to freeze some part of the variable in tensorflow 梯度 ae 乘以掩码，这不是想要的解决方案。