TensorFlow 2.5中Conv-6 CNN特定权重冻结后训练时非零权重数量异常增加的问题排查

你之前的思路里，tf.stop_gradient只是在单次张量操作中临时阻止了梯度传递，但当你把处理后的张量赋值回模型的可训练变量后，优化器在训练过程中依然会对整个变量（包括那些被你设为0的权重）计算并应用梯度。这就导致原本冻结的0值权重被更新，最终非零数量增加，达不到你想要的效果。

下面提供三种不同的实现方式，你可以根据自己的使用习惯选择：

方式一：自定义训练循环（最灵活可控） #

如果你愿意脱离model.fit()的封装，自定义训练循环能让你完全掌控梯度的计算和更新过程，这是最可靠的方式：

# 先准备好优化器和损失函数
optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)
loss_fn = tf.keras.losses.CategoricalCrossentropy(from_logits=False)

# 预先生成掩码：1代表需要冻结的权重位置（原权重<0.1的部分），0代表可训练
conv1_mask = tf.cast(best_model.trainable_weights[0] < 0.1, tf.float32)
conv6_mask = tf.cast(best_model.trainable_weights[10] < 0.1, tf.float32)

# 先把初始的0值权重赋值给模型
pruned_model.trainable_weights[0].assign(tf.where(conv1_mask == 1, 0., best_model.trainable_weights[0]))
pruned_model.trainable_weights[10].assign(tf.where(conv6_mask == 1, 0., best_model.trainable_weights[10]))

# 定义训练步骤的函数（用@tf.function加速）
@tf.function
def train_step(x_batch, y_batch):
    with tf.GradientTape() as tape:
        preds = pruned_model(x_batch, training=True)
        loss = loss_fn(y_batch, preds)
    
    # 计算所有可训练变量的梯度
    grads = tape.gradient(loss, pruned_model.trainable_weights)
    
    # 对需要冻结的权重，把对应梯度置为0——这样优化器就不会更新它们
    grads[0] = grads[0] * (1 - conv1_mask)
    grads[10] = grads[10] * (1 - conv6_mask)
    
    # 应用梯度更新
    optimizer.apply_gradients(zip(grads, pruned_model.trainable_weights))
    
    # 计算当前批次的准确率
    acc = tf.reduce_mean(tf.keras.metrics.categorical_accuracy(y_batch, preds))
    return loss, acc

# 开始训练循环
for epoch in range(10):
    print(f"Epoch {epoch+1}/10")
    total_loss = 0.0
    total_acc = 0.0
    batch_count = len(X_train) // 32  # 假设用32的批次大小
    
    for batch_idx in range(batch_count):
        x_batch = X_train[batch_idx*32 : (batch_idx+1)*32]
        y_batch = y_train[batch_idx*32 : (batch_idx+1)*32]
        loss, acc = train_step(x_batch, y_batch)
        total_loss += loss
        total_acc += acc
    
    # 打印本轮训练的平均指标
    avg_loss = total_loss / batch_count
    avg_acc = total_acc / batch_count
    val_loss, val_acc = pruned_model.evaluate(X_test, y_test, verbose=0)
    print(f"Train Loss: {avg_loss:.4f} | Train Acc: {avg_acc:.4f} | Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.4f}")

方式二：自定义带掩码的卷积层（适配model.fit()） #

如果你想继续使用model.fit()，可以把需要冻结部分权重的卷积层替换成自定义层，让层自动处理掩码和梯度过滤：

class MaskedConv2D(tf.keras.layers.Conv2D):
    def __init__(self, mask, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.mask = mask  # 掩码：1=冻结，0=可训练
    
    def call(self, inputs):
        # 每次前向传播时，确保冻结的权重保持为0
        self.kernel.assign(self.kernel * (1 - self.mask))
        return super().call(inputs)
    
    def get_gradients(self, loss, inputs):
        # 获取梯度后，把冻结位置的梯度置为0
        grads = super().get_gradients(loss, inputs)
        grads[0] = grads[0] * (1 - self.mask)
        return grads

# 重新构建模型，替换第一层和第六层卷积层
def conv6_cnn_with_mask(conv1_mask, conv6_mask):
    model = tf.keras.Sequential([
        # 第一层替换为带掩码的卷积层
        MaskedConv2D(
            mask=conv1_mask,
            filters=64, kernel_size=(3,3), activation='relu',
            kernel_initializer=tf.initializers.GlorotNormal(),
            strides=(1,1), padding='same', input_shape=(32,32,3)
        ),
        tf.keras.layers.Conv2D(64, (3,3), activation='relu', padding='same'),
        tf.keras.layers.MaxPooling2D((2,2)),
        tf.keras.layers.Conv2D(128, (3,3), activation='relu', padding='same'),
        tf.keras.layers.Conv2D(128, (3,3), activation='relu', padding='same'),
        tf.keras.layers.MaxPooling2D((2,2)),
        tf.keras.layers.Conv2D(256, (3,3), activation='relu', padding='same'),
        # 第六层替换为带掩码的卷积层
        MaskedConv2D(
            mask=conv6_mask,
            filters=256, kernel_size=(3,3), activation='relu',
            kernel_initializer=tf.initializers.GlorotNormal(),
            strides=(1,1), padding='same'
        ),
        tf.keras.layers.MaxPooling2D((2,2)),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(256, activation='relu'),
        tf.keras.layers.Dense(256, activation='relu'),
        tf.keras.layers.Dense(10, activation='softmax')
    ])
    
    model.compile(
        loss=tf.keras.losses.CategoricalCrossentropy(from_logits=False),
        optimizer=tf.keras.optimizers.Adam(learning_rate=0.01),
        metrics=['accuracy']
    )
    return model

# 生成掩码并构建模型
conv1_mask = tf.cast(best_model.trainable_weights[0] < 0.1, tf.float32)
conv6_mask = tf.cast(best_model.trainable_weights[10] < 0.1, tf.float32)

pruned_model = conv6_cnn_with_mask(conv1_mask, conv6_mask)
pruned_model.set_weights(best_model.get_weights())

# 初始化冻结的权重为0
pruned_model.layers[0].kernel.assign(pruned_model.layers[0].kernel * (1 - conv1_mask))
pruned_model.layers[6].kernel.assign(pruned_model.layers[6].kernel * (1 - conv6_mask))

# 正常训练即可
history = pruned_model.fit(
    X_train, y_train,
    epochs=10,
    validation_data=(X_test, y_test)
)

方式三：使用训练回调（最简单的适配model.fit()） #

如果你不想修改模型结构，可以创建一个回调函数，在每次训练批次前后强制把冻结的权重重置为0：

class MaskedWeightCallback(tf.keras.callbacks.Callback):
    def __init__(self, weight_indices, masks):
        super().__init__()
        self.weight_indices = weight_indices  # 需要处理的权重索引
        self.masks = masks  # 对应的掩码
    
    def on_train_batch_begin(self, batch, logs=None):
        # 批次训练前，确保冻结权重为0
        for idx, mask in zip(self.weight_indices, self.masks):
            var = self.model.trainable_weights[idx]
            var.assign(var * (1 - mask))
    
    def on_train_batch_end(self, batch, logs=None):
        # 批次训练后，再次重置冻结权重为0（防止梯度更新改变它们）
        for idx, mask in zip(self.weight_indices, self.masks):
            var = self.model.trainable_weights[idx]
            var.assign(var * (1 - mask))

# 生成掩码
conv1_mask = tf.cast(best_model.trainable_weights[0] < 0.1, tf.float32)
conv6_mask = tf.cast(best_model.trainable_weights[10] < 0.1, tf.float32)

# 初始化模型并加载权重
pruned_model = conv6_cnn()
pruned_model.set_weights(best_model.get_weights())

# 先把冻结的权重设为0
pruned_model.trainable_weights[0].assign(pruned_model.trainable_weights[0] * (1 - conv1_mask))
pruned_model.trainable_weights[10].assign(pruned_model.trainable_weights[10] * (1 - conv6_mask))

# 编译模型
pruned_model.compile(
    loss=tf.keras.losses.CategoricalCrossentropy(from_logits=False),
    optimizer=tf.keras.optimizers.Adam(learning_rate=0.01),
    metrics=['accuracy']
)

# 创建回调并开始训练
mask_callback = MaskedWeightCallback(
    weight_indices=[0, 10],
    masks=[conv1_mask, conv6_mask]
)

history = pruned_model.fit(
    X_train, y_train,
    epochs=10,
    validation_data=(X_test, y_test),
    callbacks=[mask_callback]
)

• 核心逻辑是：要么阻止优化器对冻结权重计算梯度（自定义循环/自定义层），要么在每次训练前后强制重置冻结权重为0（回调方法）。
• 你之前用tf.stop_gradient的思路只在单次操作中有效，没有从根本上阻止优化器更新这些权重，所以才会出现非零数量增加的问题。

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