如何在TensorFlow网络中停用权重？实现70%连接裁剪

嘿，这个问题其实核心就是把那些被置0的权重的梯度锁死，不让优化器对它们进行更新。下面我给你几个实用的实现方案，不管你用PyTorch还是TensorFlow都能直接套用：

核心思路 #

我们需要一个固定的**掩码（mask）**数组（0表示切断连接，1表示保留连接），然后确保：

1. 初始权重和掩码相乘，先切断70%的连接；
2. 每次反向传播后，把权重的梯度和掩码相乘，让切断位置的梯度为0，这样优化器就不会更新这些权重。

方案1：PyTorch手动处理梯度（最直观） #

先定义模型并初始化掩码，然后在训练循环里手动掩码梯度：

import torch
import torch.nn as nn
import torch.optim as optim

# 定义两层全连接网络
class SparseTwoLayerNet(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, output_dim)
        
        # 创建固定掩码：随机保留30%的连接（70%切断）
        self.mask_fc1 = torch.rand_like(self.fc1.weight) > 0.7  # 布尔数组，True对应保留
        self.mask_fc2 = torch.rand_like(self.fc2.weight) > 0.7
        
        # 初始切断连接：权重和掩码相乘
        self.fc1.weight.data *= self.mask_fc1.float()
        self.fc2.weight.data *= self.mask_fc2.float()

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 初始化模型和优化器
model = SparseTwoLayerNet(input_dim=10, hidden_dim=20, output_dim=2)
optimizer = optim.SGD(model.parameters(), lr=0.01)
loss_fn = nn.CrossEntropyLoss()

# 训练循环
for epoch in range(100):
    optimizer.zero_grad()
    
    # 前向传播
    inputs = torch.randn(32, 10)  # 示例输入
    targets = torch.randint(0, 2, (32,))
    outputs = model(inputs)
    
    # 计算损失+反向传播
    loss = loss_fn(outputs, targets)
    loss.backward()
    
    # 关键步骤：用掩码冻结切断位置的梯度
    with torch.no_grad():
        model.fc1.weight.grad *= model.mask_fc1.float()
        model.fc2.weight.grad *= model.mask_fc2.float()
    
    # 更新权重（只有保留的连接会被更新）
    optimizer.step()

方案2：PyTorch用梯度钩子（更简洁） #

可以给权重注册一个梯度钩子，自动在反向传播时处理梯度，不用每次在训练循环里写重复代码：

class SparseTwoLayerNet(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, output_dim)
        
        # 创建掩码并初始化权重
        self.mask_fc1 = torch.rand_like(self.fc1.weight) > 0.7
        self.mask_fc2 = torch.rand_like(self.fc2.weight) > 0.7
        self.fc1.weight.data *= self.mask_fc1.float()
        self.fc2.weight.data *= self.mask_fc2.float()
        
        # 注册梯度钩子：自动掩码梯度
        self.fc1.weight.register_hook(lambda grad: grad * self.mask_fc1.float())
        self.fc2.weight.register_hook(lambda grad: grad * self.mask_fc2.float())

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 训练循环就不用手动处理梯度了，和普通训练一样
for epoch in range(100):
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = loss_fn(outputs, targets)
    loss.backward()
    optimizer.step()

方案3：TensorFlow实现 #

TensorFlow里的思路完全一致，用GradientTape获取梯度后手动掩码：

import tensorflow as tf

# 定义模型
input_dim = 10
hidden_dim = 20
output_dim = 2

inputs = tf.keras.Input(shape=(input_dim,))
x = tf.keras.layers.Dense(hidden_dim, activation='relu')(inputs)
outputs = tf.keras.layers.Dense(output_dim)(x)
model = tf.keras.Model(inputs=inputs, outputs=outputs)

# 创建固定掩码
mask_fc1 = tf.random.uniform(shape=model.layers[1].kernel.shape) > 0.7
mask_fc1 = tf.cast(mask_fc1, tf.float32)
mask_fc2 = tf.random.uniform(shape=model.layers[2].kernel.shape) > 0.7
mask_fc2 = tf.cast(mask_fc2, tf.float32)

# 初始切断连接
model.layers[1].kernel.assign(model.layers[1].kernel * mask_fc1)
model.layers[2].kernel.assign(model.layers[2].kernel * mask_fc2)

# 自定义训练步骤
optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

@tf.function
def train_step(inputs, targets):
    with tf.GradientTape() as tape:
        predictions = model(inputs, training=True)
        loss = loss_fn(targets, predictions)
    
    # 获取梯度并应用掩码
    grads = tape.gradient(loss, model.trainable_variables)
    grads[0] *= mask_fc1  # 第一个可训练变量是fc1的kernel
    grads[2] *= mask_fc2  # 第三个是fc2的kernel（第二个是fc1的bias）
    
    # 更新权重
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss

# 训练循环
for epoch in range(100):
    inputs = tf.random.normal((32, input_dim))
    targets = tf.random.uniform((32,), maxval=output_dim, dtype=tf.int32)
    loss = train_step(inputs, targets)

关键注意事项 #

• 掩码要固定不变：训练过程中不要重新生成掩码，不然切断的连接会随机变化，达不到固定稀疏结构的目的。
• bias的处理：如果你也想切断bias的连接，同理给bias创建掩码即可（不过一般bias都是保留的，根据你的需求调整）。
• 稀疏性验证：训练几轮后可以打印权重的非零比例，确认是不是保持在30%左右，确保逻辑正确。

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