tf.keras.ops.custom_gradient

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

tf.keras.ops.custom_gradient(
    f
)

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

Примеры:

1. Пример, не зависящий от backend.

@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 — в зависимости от выбранного backend. Для backend JAX и TensorFlow требуется только один аргумент, тогда как для backend PyTorch может использоваться аргумент upstream.

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

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