可学习的 LeakyReLU 激活 function 和 Pytorch

Question

I'm trying to write a class for Invertible trainable LeakyReLu in which the model modifies the negative_slope in each iteration,

class InvertibleLeakyReLU(nn.Module):
  def __init__(self, negative_slope):
    super(InvertibleLeakyReLU, self).__init__()
    self.negative_slope = torch.tensor(negative_slope, requires_grad=True)
  def forward(self, input, logdet = 0, reverse = False):
    if reverse == True:
      input = torch.where(input>=0.0, input, input *(1/self.negative_slope))

      log = - torch.where(input >= 0.0, torch.zeros_like(input), torch.ones_like(input) * math.log(self.negative_slope))
      logdet = (sum(log, dim=[1, 2, 3]) +logdet).mean()
      return input, logdet
    else:
      input = torch.where(input>=0.0, input, input *(self.negative_slope))

      log = torch.where(input >= 0.0, torch.zeros_like(input), torch.ones_like(input) * math.log(self.negative_slope)) 
      logdet = (sum(log, dim=[1, 2, 3]) +logdet).mean()
      return input, logdet 

However I set requires_grad=True , the negative slope wouldn't update. Are there any other points that I must modify?

Answer 1

Does your optimizer know it should update InvertibleLeakyReLU.negative_slope ?
My guess is - no:
self.negative_slope is not defined as nn.Parameter , and therefore, by default, when you initialize your optimizer with model.parameters() negative_slope is not one of the optimization parameters.

You can either define negative_slope as a nn.Parameter :

self.negative_slope = nn.Parameter(data=torch.tensor(negative_slope), requires_grad=True)

Or, explicitly pass negative_slope from all InvertibleLeakyReLU in your model to the optimizer.