在pytorch的神经网络中将参数约束为-1、0或1

Posted 2023-03-12

【中文标题】在pytorch的神经网络中将参数约束为-1、0或1

【英文标题】：Constrain parameters to be -1, 0 or 1 in neural network in pytorch

【发布时间】：2021-08-18 16:33:12

【问题描述】：

我想约束神经网络中中间层的参数，使其更喜欢离散值：-1、0 或 1。这个想法是添加一个自定义目标函数，如果参数采用任何其他值，则会增加损失价值。请注意，我想约束特定层的参数，而不是所有层。

如何在 pytorch 中实现这一点？我想将此自定义损失添加到训练循环中的总损失中，如下所示：

custom_loss = constrain_parameters_to_be_discrete 
loss = other_loss + custom_loss 

可能使用 Dirichlet 先验可能会有所帮助，任何指向此的指针？

【参考方案1】：

在@Shai answer 上进行扩展并将其与this answer 混合使用可以更简单地通过自定义层来实现，您可以将特定层传递到该层中。

首先，从WolframAlpha（检查here）中计算得到的torch.abs(x**2 - torch.abs(x)) 的导数将放在regularize 函数中。

现在是Constrainer 层：

class Constrainer(torch.nn.Module):
    def __init__(self, module, weight_decay=1.0):
        super().__init__()
        self.module = module
        self.weight_decay = weight_decay

        # Backward hook is registered on the specified module
        self.hook = self.module.register_full_backward_hook(self._weight_decay_hook)

    # Not working with grad accumulation, check original answer and pointers there
    # If that's needed
    def _weight_decay_hook(self, *_):
        for parameter in self.module.parameters():
            parameter.grad = self.regularize(parameter)

    def regularize(self, parameter):
        # Derivative of the regularization term created by @Shia
        sgn = torch.sign(parameter)
        return self.weight_decay * (
            (sgn - 2 * parameter) * torch.sign(1 - parameter * sgn)
        )

    def forward(self, *args, **kwargs):
        # Simply forward and args and kwargs to module
        return self.module(*args, **kwargs)

使用非常简单（如果您需要对参数施加更多/更少的力，请使用您指定的weight_decay 超参数）：

constrained_layer = Constrainer(torch.nn.Linear(20, 10), weight_decay=0.1)

现在你不用担心不同的损失函数，可以正常使用你的模型了。

【讨论】：

喜欢衍生品 ;)

【参考方案2】：

你可以使用损失函数：

def custom_loss_function(x):
  loss = torch.abs(x**2 - torch.abs(x))
  return loss.mean()

此图绘制了单个元素的建议损失：

如您所见，x=-1, 0, 1 的建议损失为零，否则为正。

请注意，如果您想将此损失应用于特定层的权重，那么您的x 是权重，而不是层的激活。

【讨论】：

@dennlinger 您需要将其应用到 parameters 而不是 activations

@Shai 可能值得一提，因为现在看起来x 是某种输出，其命名包含loss，IMO 有点混乱。

@Rakib 在这个之后你有批处理规范层吗？你如何初始化这一层的权重？您是否有可能以具有小方差的正态分布进行初始化，以使所有权重一开始都非常接近于零？您需要重新考虑在这种情况下初始化权重的方式

那你期待什么？权重初始值非常接近于零，您需要付出很大的代价才能在 -/+0.5 左右“越过障碍”。

将 lambda 设置为高值会使将值从零更改为 -/+1 变得更加困难。尝试将权重初始化为 -1, 0, 1 - 看看它如何影响收敛