На этой странице предоставлен путь для пользователей, которые хотят обучать модели в JAX и развертывать их на мобильных устройствах для получения логических выводов. Методы в этом руководстве создают tflite_model , который можно использовать непосредственно с примером кода интерпретатора LiteRT или сохранить в файл tflite FlatBuffer.
Сквозной пример см. в кратком руководстве .
Предварительное условие
Рекомендуется попробовать эту функцию с новейшим ночным пакетом Python TensorFlow.
pip install tf-nightly --upgrade
Мы будем использовать библиотеку экспорта Orbax для экспорта моделей JAX. Убедитесь, что ваша версия JAX не ниже 0.4.20 или выше.
pip install jax --upgrade
pip install orbax-export --upgrade
Преобразование моделей JAX в LiteRT
Мы используем TensorFlow SavedModel в качестве промежуточного формата между JAX и LiteRT. Если у вас есть SavedModel, для завершения процесса преобразования можно использовать существующие API-интерфейсы LiteRT.
# This code snippet converts a JAX model to TFLite through TF SavedModel.
from orbax.export import ExportManager
from orbax.export import JaxModule
from orbax.export import ServingConfig
import tensorflow as tf
import jax.numpy as jnp
def model_fn(_, x):
return jnp.sin(jnp.cos(x))
jax_module = JaxModule({}, model_fn, input_polymorphic_shape='b, ...')
# Option 1: Simply save the model via `tf.saved_model.save` if no need for pre/post
# processing.
tf.saved_model.save(
jax_module,
'/some/directory',
signatures=jax_module.methods[JaxModule.DEFAULT_METHOD_KEY].get_concrete_function(
tf.TensorSpec(shape=(None,), dtype=tf.float32, name="input")
),
options=tf.saved_model.SaveOptions(experimental_custom_gradients=True),
)
converter = tf.lite.TFLiteConverter.from_saved_model('/some/directory')
tflite_model = converter.convert()
# Option 2: Define pre/post processing TF functions (e.g. (de)?tokenize).
serving_config = ServingConfig(
'Serving_default',
# Corresponds to the input signature of `tf_preprocessor`
input_signature=[tf.TensorSpec(shape=(None,), dtype=tf.float32, name='input')],
tf_preprocessor=lambda x: x,
tf_postprocessor=lambda out: {'output': out}
)
export_mgr = ExportManager(jax_module, [serving_config])
export_mgr.save('/some/directory')
converter = tf.lite.TFLiteConverter.from_saved_model('/some/directory')
tflite_model = converter.convert()
# Option 3: Convert from TF concrete function directly
converter = tf.lite.TFLiteConverter.from_concrete_functions(
[
jax_module.methods[JaxModule.DEFAULT_METHOD_KEY].get_concrete_function(
tf.TensorSpec(shape=(None,), dtype=tf.float32, name="input")
)
]
)
tflite_model = converter.convert()
Проверьте конвертированную модель TFLite
После преобразования модели в TFLite вы можете запустить API-интерфейсы интерпретатора TFLite для проверки выходных данных модели.
# Run the model with LiteRT
interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors() input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
interpreter.set_tensor(input_details[0]["index"], input_data)
interpreter.invoke()
result = interpreter.get_tensor(output_details[0]["index"])
На этой странице предоставлен путь для пользователей, которые хотят обучать модели в JAX и развертывать их на мобильных устройствах для получения логических выводов. Методы в этом руководстве создают tflite_model , который можно использовать непосредственно с примером кода интерпретатора LiteRT или сохранить в файл tflite FlatBuffer.
Сквозной пример см. в кратком руководстве .
Предварительное условие
Рекомендуется попробовать эту функцию с новейшим ночным пакетом Python TensorFlow.
pip install tf-nightly --upgrade
Мы будем использовать библиотеку экспорта Orbax для экспорта моделей JAX. Убедитесь, что ваша версия JAX не ниже 0.4.20 или выше.
pip install jax --upgrade
pip install orbax-export --upgrade
Преобразование моделей JAX в LiteRT
Мы используем TensorFlow SavedModel в качестве промежуточного формата между JAX и LiteRT. Если у вас есть SavedModel, для завершения процесса преобразования можно использовать существующие API-интерфейсы LiteRT.
# This code snippet converts a JAX model to TFLite through TF SavedModel.
from orbax.export import ExportManager
from orbax.export import JaxModule
from orbax.export import ServingConfig
import tensorflow as tf
import jax.numpy as jnp
def model_fn(_, x):
return jnp.sin(jnp.cos(x))
jax_module = JaxModule({}, model_fn, input_polymorphic_shape='b, ...')
# Option 1: Simply save the model via `tf.saved_model.save` if no need for pre/post
# processing.
tf.saved_model.save(
jax_module,
'/some/directory',
signatures=jax_module.methods[JaxModule.DEFAULT_METHOD_KEY].get_concrete_function(
tf.TensorSpec(shape=(None,), dtype=tf.float32, name="input")
),
options=tf.saved_model.SaveOptions(experimental_custom_gradients=True),
)
converter = tf.lite.TFLiteConverter.from_saved_model('/some/directory')
tflite_model = converter.convert()
# Option 2: Define pre/post processing TF functions (e.g. (de)?tokenize).
serving_config = ServingConfig(
'Serving_default',
# Corresponds to the input signature of `tf_preprocessor`
input_signature=[tf.TensorSpec(shape=(None,), dtype=tf.float32, name='input')],
tf_preprocessor=lambda x: x,
tf_postprocessor=lambda out: {'output': out}
)
export_mgr = ExportManager(jax_module, [serving_config])
export_mgr.save('/some/directory')
converter = tf.lite.TFLiteConverter.from_saved_model('/some/directory')
tflite_model = converter.convert()
# Option 3: Convert from TF concrete function directly
converter = tf.lite.TFLiteConverter.from_concrete_functions(
[
jax_module.methods[JaxModule.DEFAULT_METHOD_KEY].get_concrete_function(
tf.TensorSpec(shape=(None,), dtype=tf.float32, name="input")
)
]
)
tflite_model = converter.convert()
Проверьте конвертированную модель TFLite
После преобразования модели в TFLite вы можете запустить API-интерфейсы интерпретатора TFLite для проверки выходных данных модели.
# Run the model with LiteRT
interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors() input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
interpreter.set_tensor(input_details[0]["index"], input_data)
interpreter.invoke()
result = interpreter.get_tensor(output_details[0]["index"])
На этой странице предоставлен путь для пользователей, которые хотят обучать модели в JAX и развертывать их на мобильных устройствах для получения логических выводов. Методы в этом руководстве создают tflite_model , который можно использовать непосредственно с примером кода интерпретатора LiteRT или сохранить в файл tflite FlatBuffer.
Сквозной пример см. в кратком руководстве .
Предварительное условие
Рекомендуется попробовать эту функцию с новейшим ночным пакетом Python TensorFlow.
pip install tf-nightly --upgrade
Мы будем использовать библиотеку экспорта Orbax для экспорта моделей JAX. Убедитесь, что ваша версия JAX не ниже 0.4.20 или выше.
pip install jax --upgrade
pip install orbax-export --upgrade
Преобразование моделей JAX в LiteRT
Мы используем TensorFlow SavedModel в качестве промежуточного формата между JAX и LiteRT. Если у вас есть SavedModel, для завершения процесса преобразования можно использовать существующие API-интерфейсы LiteRT.
# This code snippet converts a JAX model to TFLite through TF SavedModel.
from orbax.export import ExportManager
from orbax.export import JaxModule
from orbax.export import ServingConfig
import tensorflow as tf
import jax.numpy as jnp
def model_fn(_, x):
return jnp.sin(jnp.cos(x))
jax_module = JaxModule({}, model_fn, input_polymorphic_shape='b, ...')
# Option 1: Simply save the model via `tf.saved_model.save` if no need for pre/post
# processing.
tf.saved_model.save(
jax_module,
'/some/directory',
signatures=jax_module.methods[JaxModule.DEFAULT_METHOD_KEY].get_concrete_function(
tf.TensorSpec(shape=(None,), dtype=tf.float32, name="input")
),
options=tf.saved_model.SaveOptions(experimental_custom_gradients=True),
)
converter = tf.lite.TFLiteConverter.from_saved_model('/some/directory')
tflite_model = converter.convert()
# Option 2: Define pre/post processing TF functions (e.g. (de)?tokenize).
serving_config = ServingConfig(
'Serving_default',
# Corresponds to the input signature of `tf_preprocessor`
input_signature=[tf.TensorSpec(shape=(None,), dtype=tf.float32, name='input')],
tf_preprocessor=lambda x: x,
tf_postprocessor=lambda out: {'output': out}
)
export_mgr = ExportManager(jax_module, [serving_config])
export_mgr.save('/some/directory')
converter = tf.lite.TFLiteConverter.from_saved_model('/some/directory')
tflite_model = converter.convert()
# Option 3: Convert from TF concrete function directly
converter = tf.lite.TFLiteConverter.from_concrete_functions(
[
jax_module.methods[JaxModule.DEFAULT_METHOD_KEY].get_concrete_function(
tf.TensorSpec(shape=(None,), dtype=tf.float32, name="input")
)
]
)
tflite_model = converter.convert()
Проверьте конвертированную модель TFLite
После преобразования модели в TFLite вы можете запустить API-интерфейсы интерпретатора TFLite для проверки выходных данных модели.
# Run the model with LiteRT
interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors() input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
interpreter.set_tensor(input_details[0]["index"], input_data)
interpreter.invoke()
result = interpreter.get_tensor(output_details[0]["index"])