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

Question

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

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

custom_loss = constrain_parameters_to_be_discrete 
loss = other_loss + custom_loss 

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

Answer 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)

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

Answer 2

你可以使用损失函数：

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

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

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

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