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

Question

我想将神经网络中的中间层的参数限制为更喜欢离散值:-1、0 或 1。这个想法是添加一个自定义目标函数，如果参数采用任何其他值，则会增加损失。请注意，我想约束特定层的参数，而不是所有层。
我如何在 pytorch 中实现它？我想将此自定义损失添加到训练循环中的总损失中，如下所示:

custom_loss = constrain_parameters_to_be_discrete 
loss = other_loss + custom_loss 

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

Answer 1

延伸至 @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)

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