如何创建带有条件的 PyTorch 钩子？答案

【问题标题】：How to create a PyTorch hook with conditions?如何创建带有条件的 PyTorch 钩子？
【发布时间】：2021-12-17 09:37:21
【问题描述】：

我正在学习钩子并使用二值化神经网络。问题是有时我的梯度在反向传递中为 0。我正在尝试用某个值替换这些渐变。

假设我有以下网络

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

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.fc1 = nn.Linear(1, 2)
        self.fc2 = nn.Linear(2, 3)
        self.fc3 = nn.Linear(3, 1)

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

net = Model()

opt = optim.Adam(net.parameters())

还有一些功能

features = torch.rand((3,1))

我可以正常训练它：

for i in range(10):
    opt.zero_grad()
    out = net(features)
    loss = torch.mean(torch.square(torch.tensor(5) - torch.sum(out)))
    loss.backward()
    opt.step()

如何附加一个钩子函数，该函数将具有以下条件用于向后传递（对于每一层）：

如果单层的所有梯度都为0，则改为1.0。
如果其中一个梯度为 0，但至少有一个梯度不为 0，则将其更改为 0.5。

【问题讨论】：

标签： python machine-learning pytorch backpropagation

【解决方案1】：

您可以使用nn.Module.register_full_backward_hook 在nn.Module 上附加回调函数：

您必须处理这两种情况：如果所有元素都等于 0，则使用 torch.all，否则（即至少一个非零）如果至少一个元素等于 0，则使用 @ 987654323@.

def grad_mod(module, grad_inputs, grad_outputs):
    if module.weight.grad is None: # safety measure for last layer 
        return None                # and layers w/ require_grad=False

    flat = module.weight.grad.view(-1)
    if torch.all(flat == 0):
        flat.data.fill_(1.)
    elif torch.any(flat == 0):
        flat.data.scatter_(0, (flat == 0).nonzero()[:,0], value=.5)

第一个子句中的指令会将所有值填充到1.，而第二个子句中的指令只会将零值替换为.5。

将挂钩连接到nn.Module：

>>> net.fc3.register_full_backward_hook(grad_mod)

这里我使用print突变flat前后的语句来展示钩子的效果：

>>> net(torch.rand((3,1))).backward(torch.tensor([[0],[1],[2]]))
>>> tensor([0.0947, 0.0000, 0.0000]) # before
>>> tensor([0.0947, 0.5000, 0.5000]) # after

>>> net(torch.rand((3,1))).backward(torch.tensor([[0],[1],[2]]))
>>> tensor([0., 0., 0.])             # before
>>> tensor([1., 1., 1.])             # after

为了将此钩子应用于多个层，您可以包装 grad_mod 并利用 nn.Module.apply 递归行为：

>>> def apply_grad_mod(module):
...     if hasattr(module, 'weight'):
...         module.register_full_backward_hook(grad_mod)

然后下面将在所有层权重上应用钩子。

>>> net.apply(apply_grad_mod)

注意：如果您还希望影响偏差，则必须扩展此行为！

【讨论】：