如何在Keras或tf.keras中冻结权重矩阵内的指定权重？

Question #

I'm running an experiment where I need to freeze specific individual weights in my Keras model. Here's my sample code:

from keras.models import Sequential
from keras.layers import Dense,Input
model = Sequential()
model.add(Dense(4, input_shape=(4,),activation='relu'))
model.add(Dense(3,name="hidden",activation='relu'))
model.add(Dense(2,activation='sigmoid'))
model.compile(loss='mse', optimizer='adam')
print(model.layers[1].get_weights()[0])

This code prints the weight matrix from the input layer to the hidden layer:

array([[-0.05557871, 0.10941017, -0.59108734],
       [ 0.37056673, 0.2968588 , 0.14362943],
       [-0.05471832, -0.21425706, 0.6455065 ],
       [-0.7883829 , -0.23868048, -0.517396 ]], dtype=float32)

I want to freeze the weights at 2nd row, 3rd column (value=0.14362943) and 4th row, 2nd column (value=-0.23868048) so they don't update during backpropagation. How can I achieve this?

Answer #

Great question! Freezing individual weight elements (instead of entire layers) isn't supported out-of-the-box with Keras' standard fit() method, but we can implement this by customizing the training loop and manually zeroing out gradients for the specific positions during backpropagation. Here's a step-by-step solution for your case:

1. Key Concept #

We'll use TensorFlow's GradientTape to track gradients, then modify the gradient tensor for our target weights—setting the gradients at the positions we want to freeze to 0. This ensures those weights won't be updated when we apply the optimizer.

2. Full Implementation Code #

import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense

# Build your model (same as before)
model = Sequential()
model.add(Dense(4, input_shape=(4,), activation='relu'))
model.add(Dense(3, name="hidden", activation='relu'))
model.add(Dense(2, activation='sigmoid'))

# Define optimizer and loss function (instead of using model.compile())
optimizer = tf.keras.optimizers.Adam()
loss_fn = tf.keras.losses.MeanSquaredError()

# Define the positions to freeze (note: Python uses 0-based indexing)
# 2nd row, 3rd column → row index 1, column index 2
# 4th row, 2nd column → row index 3, column index 1
freeze_positions = [(1, 2), (3, 1)]

# Get the target weight tensor (kernel of the hidden layer)
target_weights = model.get_layer("hidden").kernel

# Custom training loop
def train_step(x, y):
    with tf.GradientTape() as tape:
        # Forward pass
        predictions = model(x, training=True)
        # Calculate loss
        loss = loss_fn(y, predictions)
    
    # Get gradients for all trainable variables
    grads = tape.gradient(loss, model.trainable_variables)
    
    # Find the index of our target weights in the trainable variables list
    target_idx = model.trainable_variables.index(target_weights)
    target_grads = grads[target_idx]
    
    # Zero out gradients at the specified positions
    for (row, col) in freeze_positions:
        target_grads = tf.tensor_scatter_nd_update(
            target_grads,
            indices=[[row, col]],
            updates=[0.0]
        )
    
    # Update the gradients list with modified gradients
    grads[target_idx] = target_grads
    
    # Apply gradients to update weights
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    
    return loss

# Example usage (replace with your actual dataset)
import numpy as np
x_train = np.random.rand(100, 4)
y_train = np.random.rand(100, 2)

# Run training
epochs = 10
for epoch in range(epochs):
    epoch_loss = 0.0
    for x, y in zip(x_train, y_train):
        x = tf.expand_dims(x, 0)  # Add batch dimension
        y = tf.expand_dims(y, 0)
        loss = train_step(x, y)
        epoch_loss += loss.numpy()
    print(f"Epoch {epoch+1}, Loss: {epoch_loss/len(x_train):.4f}")

# Verify the frozen weights haven't changed
print("\nFrozen weights after training:")
final_weights = model.get_layer("hidden").get_weights()[0]
print(f"2nd row, 3rd column: {final_weights[1,2]}")
print(f"4th row, 2nd column: {final_weights[3,1]}")

3. Explanation #

• Custom Training Loop: We replace model.fit() with a manual loop to have full control over gradient computation.
• Gradient Modification: Using tf.tensor_scatter_nd_update, we set the gradients at our target positions to 0. Since gradients drive weight updates, a 0 gradient means no change to that weight.
• Verification: At the end, we print the frozen weights to confirm they remain unchanged after training.

This approach gives you precise control over which individual weights to freeze, even when working with standard Keras layers.

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