RuntimeError: «элемент 0 тензоров не требует градации» при использовании пользовательской функции автоградации с create_ ⇐ Python

[Anonymous]

Вот уникальный, высококачественный вопрос о машинном обучении, созданный для Stack Overflow.
Он посвящен тонкой, но сложной проблеме в PyTorch: реализации пользовательской функции autograd.Function, которая дает сбой при вычислении производных второго порядка (обычно требуется для таких методов, как градиентный штраф в WGAN, векторные продукты Гессе или мета-обучение/MAML).
Этот вопрос предназначен для «недоступно для Google», но решаемо экспертом, что обычно привлекает большое внимание.
Черновик вопроса о переполнении стека
Title: RuntimeError: «элемент 0 тензоров не требует градации» при использовании пользовательской функции autograd с create_graph=True
Теги: python pytorch autograd глубокое обучение градиентный спуск
Тело:
Я реализую специальную функцию активации (вариант Swish) в PyTorch для оптимизации использования памяти. Я реализовал его с помощью torch.autograd.Function, определив как прямые, так и обратные статические методы.
Стандартное обучение (производные первого порядка) работает отлично. Однако сейчас я пытаюсь использовать этот пользовательский слой в настройке WGAN-GP (Wasserstein GAN с градиентным штрафом). Для этого требуется вычислить градиент градиентов (двойное обратное свойство) с помощью torch.autograd.grad(..., create_graph=True).
Как только я включаю create_graph=True, я получаю следующую ошибку: RuntimeError: элемент 0 тензоров не требует grad и не имеет grad_fn.
Что я пробовал:

• I проверил, что x.requires_grad имеет значение True.
• Я попробовал заменить свою пользовательскую функцию на torch.nn.SiLU() (встроенную в
  Swish), и двойная обратная связь работает отлично, так что проблема
  определенно в моем классе CustomSwish.
• Я удалил ctx.save_for_backward и пересчитал вводит данные в обратном направлении,
  но ошибка сохраняется.

Вопрос: Почему моя пользовательская функция autograd.Function нарушает граф вычислений во время второго обратного прохода, хотя я использую дифференцируемые операции PyTorch внутри в обратном направлении? Как правильно реализовать пользовательскую функцию, поддерживающую производные более высокого порядка?

Подробнее здесь: https://stackoverflow.com/questions/798 ... custom-aut