tf.keras.ops.custom_gradient

Декоратор для определения функции с пользовательским градиентом.

tf.keras.ops.custom_gradient(
    f
)

Этот декоратор позволяет точно управлять градиентами последовательности операций. Это может быть полезно по многим причинам, включая обеспечение более эффективного или численно стабильного градиента для последовательности операций.

Примеры:

1. Примеp, не зависящий от бэкэнда.

@ops.custom_gradient
def log1pexp(x):
    e = ops.exp(x)

    def grad(*args, upstream=None):
        if upstream is None:
            (upstream,) = args
        return ops.multiply(upstream, 1.0 - 1.0 / ops.add(1, e))

    return ops.log(1 + e), grad

Обратите внимание, что функция grad, возвращающая вычисление градиента, требует args, а также аргумент upstream, в зависимости от установленного бэкэнда. С бэкендами JAX и TensorFlow требуется только один аргумент, тогда как в случае бэкэнда PyTorch может использоваться аргумент upstream.

При работе с бэкендом TensorFlow/JAX достаточно grad(upstream). С PyTorch функция grad требует *args а также upstream, например, def grad(*args, upstream). Следуйте предыдущему примеру, чтобы использовать @ops.custom_gradient способом, совместимым со всеми бэкендами.

1. Вот пример, специфичный для JAX и TensorFlow:

@ops.custom_gradient
def log1pexp(x):
    e = ops.exp(x)
    def grad(upstream):
        return ops.multiply(upstream, 1.0 - 1.0 / ops.add(1, e))
    return ops.log(1 + e), grad

1. Наконец, вот пример, специфичный для PyTorch, использующий *args и upstream:

@ops.custom_gradient
def log1pexp(x):
    e = ops.exp(x)
    def grad(*args, upstream):
        return ops.multiply(upstream, 1.0 - 1.0 / ops.add(1, e))
    return ops.log(1 + e), grad