为何用register_hook()冻结部分权重张量无法加速训练？

关于register_hook()的机制 #

register_hook()是先计算完整梯度，再对梯度进行修改（置零），所以即便你把部分梯度置零，反向传播时依然会计算所有权重的梯度——那些你想冻结的权重，梯度计算过程并没有被跳过，只是最后把梯度置零了，计算量没减少，这就是训练没加速的原因。

实现类似requires_grad=False的高效冻结方案 #

要真正跳过部分权重的梯度计算、实现训练加速，你可以把权重张量拆分成可训练部分和不可训练部分，分别处理：

• 将原权重拆分为两个张量，不可训练部分设置requires_grad=False
• 前向传播时拼接两个张量作为完整权重使用
• 优化器只传入可训练的张量参数

这样反向传播时，PyTorch会自动跳过requires_grad=False部分的梯度计算，大幅减少计算量。

修改后的代码示例 #

import torch
import torch.nn as nn
import torch.optim as optim
import time
from torch.utils.data import Dataset, DataLoader

torch.manual_seed(42)

class example_net(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, freeze_cols_start=10):
        super(example_net, self).__init__()
        self.freeze_cols_start = freeze_cols_start  # 从第10列开始冻结
        
        # 拆分fc1的权重：可训练的前10列，不可训练的剩余列
        self.fc1_trainable = nn.Parameter(torch.randn(hidden_size, freeze_cols_start))
        # 不可训练部分设置requires_grad=False
        self.fc1_frozen = nn.Parameter(torch.randn(hidden_size, input_size - freeze_cols_start), requires_grad=False)
        
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.fc3 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        # 拼接可训练和不可训练的权重，得到完整的fc1权重
        fc1_weight = torch.cat([self.fc1_trainable, self.fc1_frozen], dim=1)
        # 手动计算fc1的输出（等价于nn.Linear的前向）
        x = torch.relu(torch.matmul(x, fc1_weight.T) + self.fc1_trainable.bias)
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# dataset example:
class example_dataset(Dataset):
    def __init__(self, input_size, num_samples):
        self.input_size = input_size
        self.num_samples = num_samples
        self.data = torch.randn(num_samples, input_size)
        self.targets = torch.randint(0, 10, (num_samples,))

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        return self.data[idx], self.targets[idx]

input_size = 4096
batch_size = 64
num_samples = 10000

dataset = example_dataset(input_size, num_samples)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
criterion = nn.CrossEntropyLoss()

def train_model(net, dataloader, num_epochs=5):
    # 只传入可训练的参数
    optimizer = optim.SGD([p for p in net.parameters() if p.requires_grad], lr=0.01)
    start_time = time.time()

    for epoch in range(num_epochs):
        for inputs, targets in dataloader:
            optimizer.zero_grad()
            outputs = net(inputs)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()

    end_time = time.time()
    return end_time - start_time

num_epochs = 1

# 原模型（无冻结）
net_original = example_net(input_size, 4096, 10, freeze_cols_start=input_size)
time_original = train_model(net_original, dataloader, num_epochs=num_epochs)
print(f"原模型训练时间: {time_original:.4f} seconds")

# 冻结部分权重的模型
net_frozen = example_net(input_size, 4096, 10, freeze_cols_start=10)
time_frozen = train_model(net_frozen, dataloader, num_epochs=num_epochs)
print(f"冻结部分权重后的训练时间: {time_frozen:.4f} seconds")

额外说明 #

• 如果是预训练模型，拆分权重时直接从预训练权重中截取对应部分即可，不需要随机初始化
• 这种拆分方式完全跳过了冻结部分的梯度计算，训练速度会和真正设置requires_grad=False的单参数张量一致

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