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

问题描述

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

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

custom_loss = constrain_parameters_to_be_discrete 
loss = other_loss + custom_loss 

可能使用狄利克雷先验可能会有所帮助，有任何指示吗？

解决方法

你可以使用损失函数：

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

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

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

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

,

扩展 @Shai answer 并将其与 this answer 混合可以通过自定义层实现更简单，您可以将特定层传递到其中。

首先，torch.abs(x**2 - torch.abs(x)) 的计算导数取自 WolframAlpha（检查 here）将放置在 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)

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