Utilisation d'autres frameworks ML dans TFX

TFX en tant que plate-forme est neutre en termes de framework et peut être utilisé avec d'autres frameworks ML, par exemple JAX, scikit-learn.

Pour les développeurs de modèles, cela signifie qu'ils n'ont pas besoin de réécrire le code de leur modèle implémenté dans un autre framework ML, mais peuvent plutôt réutiliser la majeure partie du code de formation tel quel dans TFX et bénéficier d'autres fonctionnalités de TFX et du reste de l'écosystème TensorFlow. des offres.

Le SDK du pipeline TFX et la plupart des modules de TFX, par exemple l'orchestrateur de pipeline, n'ont aucune dépendance directe sur TensorFlow, mais certains aspects sont orientés vers TensorFlow, tels que les formats de données. En tenant compte des besoins d'un cadre de modélisation particulier, un pipeline TFX peut être utilisé pour entraîner des modèles dans n'importe quel autre cadre de ML basé sur Python. Cela inclut Scikit-learn, XGBoost et PyTorch, entre autres. Certaines des considérations relatives à l'utilisation des composants TFX standard avec d'autres frameworks incluent :

  • SampleGen génère tf.train.Example dans des fichiers TFRecord. Il s'agit d'une représentation générique des données d'entraînement, et les composants en aval utilisent TFXIO pour les lire sous forme de Arrow/RecordBatch en mémoire, qui peut être ensuite convertie en tf.dataset , Tensors ou d'autres formats. Des formats de charge utile/fichier autres que tf.train.Example/TFRecord sont envisagés, mais pour les utilisateurs de TFXIO, cela devrait être une boîte noire.
  • Transform peut être utilisé pour générer des exemples de formation transformés, quel que soit le framework utilisé pour la formation, mais si le format du modèle n'est pas saved_model , les utilisateurs ne pourront pas intégrer le graphique de transformation dans le modèle. Dans ce cas, la prédiction du modèle doit prendre en compte les fonctionnalités transformées au lieu des fonctionnalités brutes, et les utilisateurs peuvent exécuter la transformation comme étape de prétraitement avant d'appeler la prédiction du modèle lors de la diffusion.
  • Trainer prend en charge GenericTraining afin que les utilisateurs puissent former leurs modèles à l'aide de n'importe quel framework ML.
  • Par défaut, l'évaluateur ne prend en charge que saved_model , mais les utilisateurs peuvent fournir une FDU qui génère des prédictions pour l'évaluation du modèle.

La formation d'un modèle dans un framework non basé sur Python nécessitera d'isoler un composant de formation personnalisé dans un conteneur Docker, dans le cadre d'un pipeline qui s'exécute dans un environnement conteneurisé tel que Kubernetes.

JAX

JAX est Autograd et XLA, réunis pour une recherche en apprentissage automatique haute performance. Flax est une bibliothèque et un écosystème de réseaux neuronaux pour JAX, conçus pour la flexibilité.

Avec jax2tf , nous sommes en mesure de convertir les modèles JAX/Flax formés au format saved_model , qui peut être utilisé de manière transparente dans TFX avec une formation générique et une évaluation de modèle. Pour plus de détails, consultez cet exemple .

scikit-apprendre

Scikit-learn est une bibliothèque d'apprentissage automatique pour le langage de programmation Python. Nous avons un exemple e2e avec une formation et une évaluation personnalisées dans TFX-Addons.