如何固定PyTorch中nn.Linear层的对角线元素为零且训练不更新？

你说的每个epoch重置对角线的方法确实可行，但存在两个明显问题：一是容易遗漏操作导致bug，二是batch内权重更新后对角线已发生变化，虽epoch末重置能拉回状态，但中间会产生不必要的计算冗余，不够优雅。下面是两种更可靠高效的实现方式：

方法1：自定义Linear层，用梯度钩子阻止对角线更新 #

这种方法从梯度层面入手，让对角线元素的梯度始终为零，优化器自然不会对其进行更新，属于一劳永逸的方案。

代码示例：

import torch
import torch.nn as nn

class ZeroDiagLinear(nn.Linear):
    def __init__(self, in_features, bias=True):
        # 输入输出尺寸相同，直接将out_features设为in_features
        super().__init__(in_features, in_features, bias=bias)
        # 给权重注册梯度钩子，反向传播时自动清零对角线梯度
        self.weight.register_hook(self._zero_diag_grad)
        # 初始化阶段先把对角线置零
        with torch.no_grad():
            self.weight.diagonal().zero_()
    
    def _zero_diag_grad(self, grad):
        # 将梯度的对角线元素置零后返回
        grad.diagonal().zero_()
        return grad

使用时直接实例化该类即可，训练过程中无需额外操作，钩子会自动处理梯度，保证对角线始终为零。

方法2：优化器更新后手动清零（备选） #

如果不想自定义层，也可以在每次优化器更新参数后，手动将对角线置零。注意要放在optimizer.step()之后，并用torch.no_grad()包裹，避免干扰梯度计算：

# 初始化普通Linear层
linear_layer = nn.Linear(5, 5)
# 初始阶段先把对角线置零
with torch.no_grad():
    linear_layer.weight.diagonal().zero_()

# 训练循环内的操作
for epoch in range(num_epochs):
    for inputs, targets in dataloader:
        optimizer.zero_grad()
        output = linear_layer(inputs)
        loss = your_loss_fn(output, targets)
        loss.backward()
        optimizer.step()
        
        # 参数更新后立刻清零对角线
        with torch.no_grad():
            linear_layer.weight.diagonal().zero_()

这种方法的缺点是需要手动维护清零步骤，若存在多个同类层，代码会冗余繁琐，不如方法1省心。

另外不推荐使用forward时加掩码的方案——这种方式仅在推理阶段掩盖对角线，但权重本身仍会被优化器修改，既浪费计算资源，还可能导致训练不稳定，完全没必要。

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