プラットフォームとしての TFX はフレームワークに依存せず、他の ML フレームワーク (JAX、scikit-learn など) とともに使用できます。
モデル開発者にとって、これは、別の ML フレームワークに実装されたモデル コードを書き直す必要がなく、代わりにトレーニング コードの大部分を TFX でそのまま再利用でき、TFX やその他の TensorFlow エコシステムの他の機能の恩恵を受けることができることを意味します。オファーします。
TFX パイプライン SDK と TFX のほとんどのモジュール (パイプライン オーケストレーターなど) は、TensorFlow に直接依存していませんが、データ形式など、TensorFlow を指向しているいくつかの側面があります。特定のモデリング フレームワークのニーズを考慮すると、TFX パイプラインを使用して、他の Python ベースの ML フレームワークでモデルをトレーニングできます。これには、Scikit-learn、XGBoost、PyTorch などが含まれます。他のフレームワークで標準の TFX コンポーネントを使用する場合の考慮事項には、次のようなものがあります。
- ExampleGen は、 tf.train.Example をTFRecord ファイルに出力します。これはトレーニング データの汎用表現であり、ダウンストリーム コンポーネントはTFXIOを使用してメモリ内の Arrow/RecordBatch として読み取り、さらに
tf.dataset、Tensors、またはその他の形式に変換できます。 tf.train.Example/TFRecord 以外のペイロード/ファイル形式も検討されていますが、TFXIO ユーザーにとってはブラックボックスになるはずです。 - トレーニングにどのようなフレームワークが使用されているかに関係なく、 Transformを使用して変換されたトレーニング サンプルを生成できますが、モデル形式が
saved_modelでない場合、ユーザーは変換グラフをモデルに埋め込むことができません。その場合、モデル予測では生の特徴ではなく変換された特徴を取得する必要があり、ユーザーはサービス提供時にモデル予測を呼び出す前に前処理ステップとして変換を実行できます。 - Trainer はGenericTrainingをサポートしているため、ユーザーは任意の ML フレームワークを使用してモデルをトレーニングできます。
- エバリュエーターはデフォルトで
saved_modelのみをサポートしますが、ユーザーはモデル評価の予測を生成する UDF を提供できます。
非 Python ベースのフレームワークでモデルをトレーニングするには、Kubernetes などのコンテナー化された環境で実行されるパイプラインの一部として、カスタム トレーニング コンポーネントを Docker コンテナーに分離する必要があります。
ジャックス
JAXは Autograd と XLA であり、高性能機械学習研究のために統合されました。 Flax は、柔軟性を考慮して設計された JAX 用のニューラル ネットワーク ライブラリおよびエコシステムです。
jax2tfを使用すると、トレーニングされた JAX/Flax モデルをsaved_model形式に変換できます。これは、一般的なトレーニングとモデル評価を使用して TFX でシームレスに使用できます。詳細については、この例を確認してください。
scikit-learn
Scikit-learn は、Python プログラミング言語の機械学習ライブラリです。 TFX-Addons でカスタマイズされたトレーニングと評価を備えた e2e の例があります。