Differentiable Threhold for Conv2d output

[Ajinkya_Ambatwar]

Hi I am trying to to implement learnable thresholding for the output of conv2d layers in the CNN.
I want to make the output of conv layer to 0 for values between -threshold to +threshold. Also I want to differentiate the loss function, with respect to the threshold and update the threshold with each iteration.
I have come across .clamp and the nn.threshold method in pytorch but I am unknown about about how to use them in class definition(eg. alexnet) and will autograd support this. If not what is the way around?

Thank you!!

[KFrank]

Hi Ajinkya!

@Ajinkya_Ambatwar

I want to make the output of conv layer to 0 for values between -threshold to +threshold.

You want a differentiable, “soft” threshold function. I think that
nn.Tanhshrink gives you most of what you want. You can put in
a threshold parameter like this:

threshold * nn.functional.tanhshrink (x / threshold)

Also I want to differentiate the loss function, with respect to the threshold and update the threshold with each iteration.

I’ve never done this, but I believe that if you make your
threshold value a nn.Parameter and include it in your model,
it should work.

The method in this post:

How to train my custom parameter?

You could write your own nn.Module that includes a trainable alpha parameter: class MyLayer(nn.Module): def __init__(self, *args, **kwargs): self.alpha = nn.Parameter(torch.tensor(0.)) # Other necessary setup super().__init__(*args, **kwargs) def forward(self, x): # Necessary forward computations return output Using this MyLayer module as part of a network architecture will allow you to learn alpha along with all of the other usual parameters. …

looks plausible to me (but I can’t vouch for it).

Best.

K. Frank

[Ajinkya_Ambatwar]

Thank you Frank,
I would look into the solution you are proposing