可學習的 LeakyReLU 激活 function 和 Pytorch

Question

我正在嘗試為可逆可訓練 LeakyReLu 編寫 class ，其中 model 在每次迭代中修改negative_slope，

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 

但是我設置requires_grad=True ，負斜率不會更新。 還有其他需要修改的地方嗎？

Answer 1

您的優化器是否知道它應該更新InvertibleLeakyReLU.negative_slope ？
我的猜測是 - 不：
self.negative_slope未定義為nn.Parameter ，因此，默認情況下，當您使用model.parameters()初始化優化器時， negative_slope不是優化參數之一。

您可以將negative_slope定義為nn.Parameter ：

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

或者，將 model 中所有InvertibleLeakyReLU中的negative_slope顯式傳遞給優化器。