如何在PyTorch中创建梯度掩码并应用于2×2参数张量

1. 创建梯度掩码并注册 #

梯度掩码的作用是让指定位置的梯度归零，从而阻止对应参数更新。你需要创建一个和matrix同形状的张量，第一行全为0（屏蔽梯度），第二行全为1（保留梯度），再通过register_buffer注册——这个方法会把掩码存入模型，但不会将其当作可训练参数。

代码修改如下：

matrix = torch.nn.Parameter(torch.zeros(2, 2))
# 第一行0，第二行1的掩码张量
self.register_buffer('gradient_mask', torch.tensor([[0., 0.], [1., 1.]]))

2. 应用掩码屏蔽梯度 #

有两种简单的方式将掩码应用到参数梯度上：

方式一：使用梯度钩子（自动生效）

给matrix注册反向传播钩子，每次计算出梯度后自动用掩码相乘，直接把第一行梯度置0：

# 在模型__init__方法中添加
self.matrix.register_hook(lambda grad: grad * self.gradient_mask)

方式二：手动修改梯度（更直观）

在反向传播之后、优化器更新之前，手动将梯度与掩码相乘：

loss.backward()
# 应用掩码，把第一行梯度置0
model.matrix.grad *= model.gradient_mask
optimizer.step()

完整示例代码 #

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

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.matrix = nn.Parameter(torch.zeros(2, 2))
        self.register_buffer('gradient_mask', torch.tensor([[0., 0.], [1., 1.]]))
        # 注册钩子自动处理梯度
        self.matrix.register_hook(lambda grad: grad * self.gradient_mask)

# 测试流程
model = MyModel()
optimizer = optim.SGD(model.parameters(), lr=0.1)

for epoch in range(3):
    optimizer.zero_grad()
    # 构造简单损失：让矩阵元素尽可能接近1
    loss = torch.mean((model.matrix - 1) ** 2)
    loss.backward()
    print(f"第{epoch+1}轮梯度:\n{model.matrix.grad}")
    optimizer.step()
    print(f"第{epoch+1}轮更新后矩阵:\n{model.matrix.data}\n")

运行后你会看到：第一行参数始终保持初始的0，第二行会逐步向1逼近，符合需求。

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