pytorch 中特定参数的比例梯度？

Question

假设我有一些参数为 A、B、C 的神经网络。每当对 C 应用梯度更新时，我希望它的缩放比例与正常梯度不同（即每个梯度更新是计算梯度的 2 倍或 1/3 倍）。我希望应用于所有其他参数的渐变保持不变。我将如何在 pytorch 中解决这个问题？

Answer 1

您可以使用

grad

属性直接访问张量的梯度：

C.grad *= z

，其中

z

将是您的乘法因子。

整齐包裹的一种方法是在模块上使用带有

register_full_backward_hook

的向后挂钩。这是与您的描述相对应的最小用例：

class CustomFunction(nn.Module):
    def __init__(self):
        super().__init__()
        self.A = nn.Parameter(torch.rand(1))
        self.B = nn.Parameter(torch.rand(1))
        self.C = nn.Parameter(torch.rand(1))

    def forward(self, x):
        return (x*self.A+self.B)**self.C

让我们设置一个随机输入：

>>> torch.rand(1, requires_grad=True)

您可以定义一个钩子，当在该模块上调用向后函数时将调用该钩子。为此，您可以直接利用作为参数传递的模块的属性：

def hook(module, grad_input, grad_output):
    module.C.grad *= z

然后你用

register_full_backward_hook

:

将钩子附加到你的模块的实例上

>>> torch.manual_seed(0)
>>> F = CustomFunction()
>>> F.register_full_backward_hook(hook)

梯度是：

>>> F(x).backward()
>>> F.A.grad, F.B.grad, F.C.grad
(tensor([0.0138]), tensor([0.1044]), tensor([-0.5368]))

这可以与没有钩子的推理进行比较，如：

>>> torch.manual_seed(0)
>>> F = CustomFunction()
>>> F(x).backward()
>>> F.A.grad, F.B.grad, F.C.grad
(tensor([0.0138]), tensor([0.1044]), tensor([-0.1789]))