Поддержка нескольких фреймворков

Благодаря интеграции с Ivy, теперь Kornia может использоваться с TensorFlow, JAX и NumPy.

Это можно сделать с помощью следующих функций, которые теперь являются частью API Kornia:

• kornia.to_tensorflow()
• kornia.to_jax()
• kornia.to_numpy()

Вот пример использования Kornia с TensorFlow:

import kornia
import tensorflow as tf

tf_kornia = kornia.to_tensorflow()

rgb_image = tf.random.normal((1, 3, 224, 224))
gray_image = tf_kornia.color.rgb_to_grayscale(rgb_image)

Теперь разберёмся, что происходит здесь.

1. Транспиляция Kornia в TensorFlow

   Эта строка лениво транспилирует всё в API Kornia в TensorFlow и создаёт новый модуль для транспиленной версии Kornia. Поскольку транспиляция происходит лениво, ни одна функция или класс не будет транспилирован до тех пор, пока их не вызовут.

   tf_kornia = kornia.to_tensorflow()
   

2. Вызов функции TF Kornia

   Теперь мы можем вызывать любую функцию (или класс) Kornia с аргументами TF. Однако, эта функция будет работать очень медленно по сравнению с оригинальной функцией, так как функция транспилируется в процессе выполнения. .. code-block:: python

   rgb_image = tf.random.normal((1, 3, 224, 224)) gray_image = tf_kornia.color.rgb_to_grayscale(rgb_image) # медленно

3. Последующие вызовы функций

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

   gray_image = tf_kornia.color.rgb_to_grayscale(rgb_image)  # быстро
   

4. Транспиляции в разных сессиях Python

   Возможно, вас интересует, придется ли вам ждать завершения этих долгих начальных транспиляций каждый раз, когда вы начинаете новую сессию Python? Хорошая новость в том, что при выполнении транспиляции Ivy сохраняет сгенерированный исходный код в локальной директории, поэтому если та же транспиляция будет попытаться выполнить еще раз из той же директории, она будет сразу же извлечена и использована. Kornia может использоваться с JAX и NumPy аналогичным образом:

import kornia
import numpy as np

np_kornia = kornia.to_numpy()

rgb_image = np.random.normal(size=(1, 3, 224, 224))
gray_image = np_kornia.color.rgb_to_grayscale(rgb_image)

import kornia
import jax

jax_kornia = kornia.to_jax()

rgb_image = jax.random.normal(jax.random.key(42), shape=(1, 3, 224, 224))
gray_image = jax_kornia.color.rgb_to_grayscale(rgb_image)

Ограничения

• Преобразование Kornia в TensorFlow или JAX работает для функций, классов и обучаемых модулей; преобразование в NumPy поддерживает функции и классы, но не обучаемые модули.* Транспиляция в настоящее время не работает с пользовательскими ядрами, такими как flash attention.
• Некоторые состоятельные классы в настоящее время не могут быть транспилированы, такие как оптимизаторы (torch.optim.Adam и т.д.), тренеры и загрузчики данных.
• Совместимость с нативными компиляторами (jax.jit и tf.function) в транспилированных версиях Kornia ограничена, особенно по сравнению с torch.compile для стандартного Kornia. Улучшение совместимости с этими компиляторами является одним из ключевых направлений текущего развития Ivy.

От команды Ivy

Надеемся, вам будет полезно использование Kornia с TensorFlow, JAX и NumPy! Ivy всё ещё находится в разработке, поэтому если вы обнаружите какие-либо проблемы/баги, не стесняйтесь открыть issue в репозитории ivy. Мы также будем очень признательны, если вы поставите звезду, чтобы показать свою поддержку!

Для получения дополнительной информации о Ivy мы рекомендуем ознакомиться с нашей документацией.