Пользовательский слой с автоградом pytorch делает ошибку градиента при обучении - PullRequest
Новая
       5

Пользовательский слой с автоградом pytorch делает ошибку градиента при обучении

0 голосов
/ 24 апреля 2020

Я новичок в Pytorch и хочу создать пользовательский слой, PACT Function, "https://arxiv.org/pdf/1805.06085.pdf". Я реализовал этот слой и заменил все функции ReLU в Re sNet. Но когда я начал тренировать свою модель, возникла эта ошибка.

RuntimeError: Функция _pact_functionBackward вернула недопустимый градиент по индексу 1 - получил [], но ожидаемая форма совместима с [1]

PACT Class: 

class PACT(nn.Module):
  def __init__(self, num_bits=4):
    super(PACT, self).__init__()

    self.num_bits = num_bits
    self.alpha = nn.Parameter(torch.Tensor([5.]), requires_grad=True)

    self.pact_function = PACT_function()
    self.pact_quantize = PACT_quantize(self.alpha, self.num_bits)

  def forward(self, input):
    y = self.pact_function(input, self.alpha)
    yq = self.pact_quantize(y)

  return yq

И как пересылать и назад активацию с помощью "pact_function": 

def PACT_function():
  class _pact_function(torch.autograd.function):

    # See https://github.com/obilaniu/GradOverride/blob/master/functional.py
    # See https://arxiv.org/pdf/1805.06085.pdf

    @staticmethod
    def forward(ctx, x, alpha):
      ctx.save_for_backward(x, alpha)
      return x.clamp(min=0.).min(alpha)

    @staticmethod
    def backward(ctx, grad_output):
      x, alpha = ctx.saved_variables

      lt0      = x < 0
      gta      = x > alpha
      gi       = 1.0-lt0.float()-gta.float()

      grad_x    = grad_output*gi
      grad_alpha = torch.sum(grad_output*x.ge(alpha).float())

      return grad_x, grad_alpha   # why grad_x, None acts?? but this makes the values of alpha same in all layers!

  return _pact_function.apply

Это очень странно, потому что, если я удаляю grad_alpha, возвращается как grad_x, нет, тогда код действует. Однако, когда я печатаю значения альфа в моей модели (Re sNet), все значения одинаковы, даже если эти слои PACT являются разными слоями.

  • Epoch: 10 | LR: 0.01000 [============================ 196/196 =============== ============>] Шаг: 61мс | Всего: 22s604ms | Потеря: 0,555 | A cc: 80,896% (40448/50000)
  • PACT () Параметр, содержащий: тензор ([4.5349], device = 'cuda: 0', require_grad = True)
  • PACT ( ) Параметр, содержащий: тензор ([4.5349], устройство = 'cuda: 0', require_grad = True)
  • PACT () Параметр, содержащий: тензор ([4.5349], устройство = 'cuda: 0', require_grad = True)
  • PACT () Параметр, содержащий: tenor ([4.5349], device = 'cuda: 0', require_grad = True)
  • PACT () Параметр, содержащий: tenor ([4.5349], device = 'cuda: 0', require_grad = True)

Итак, мне интересно, почему возврат (grad_x, None) действует, а не (grad_x, grad_alpha), и почему альфа-параметры одинаковы в все слои при использовании (grad_x, None) с nn.Parameter. Очень благодарен, если у вас есть ответы.

...