O TFX como plataforma é neutro em termos de estrutura e pode ser usado com outras estruturas de ML, por exemplo, JAX, scikit-learn.
Para desenvolvedores de modelos, isso significa que eles não precisam reescrever o código do modelo implementado em outra estrutura de ML, mas podem reutilizar a maior parte do código de treinamento como está no TFX e se beneficiar de outros recursos do TFX e do restante do ecossistema TensorFlow ofertas.
O SDK do pipeline do TFX e a maioria dos módulos do TFX, por exemplo, orquestrador de pipeline, não têm nenhuma dependência direta do TensorFlow, mas há alguns aspectos orientados para o TensorFlow, como formatos de dados. Com alguma consideração das necessidades de uma estrutura de modelagem específica, um pipeline TFX pode ser usado para treinar modelos em qualquer outra estrutura de ML baseada em Python. Isso inclui Scikit-learn, XGBoost e PyTorch, entre outros. Algumas das considerações para usar os componentes padrão do TFX com outras estruturas incluem:
- ExampleGen gera tf.train.Example em arquivos TFRecord. É uma representação genérica para dados de treinamento, e os componentes downstream usam TFXIO para lê-los como Arrow/RecordBatch na memória, que pode ser posteriormente convertido em
tf.dataset,Tensorsou outros formatos. Formatos de carga/arquivo diferentes de tf.train.Example/TFRecord estão sendo considerados, mas para usuários do TFXIO deve ser uma caixa preta. - Transform pode ser usado para gerar exemplos de treinamento transformados, independentemente da estrutura usada para treinamento, mas se o formato do modelo não for
saved_model, os usuários não poderão incorporar o gráfico de transformação no modelo. Nesse caso, a previsão do modelo precisa usar recursos transformados em vez de recursos brutos, e os usuários podem executar a transformação como uma etapa de pré-processamento antes de chamar a previsão do modelo durante a veiculação. - O Trainer oferece suporte ao GenericTraining para que os usuários possam treinar seus modelos usando qualquer estrutura de ML.
- Por padrão, o avaliador oferece suporte apenas
saved_model, mas os usuários podem fornecer uma UDF que gera previsões para avaliação do modelo.
O treinamento de um modelo em uma estrutura não baseada em Python exigirá o isolamento de um componente de treinamento personalizado em um contêiner Docker, como parte de um pipeline que está sendo executado em um ambiente em contêiner, como o Kubernetes.
JAX
JAX é Autograd e XLA, reunidos para pesquisas de aprendizado de máquina de alto desempenho. Flax é uma biblioteca de rede neural e ecossistema para JAX, projetada para oferecer flexibilidade.
Com jax2tf , podemos converter modelos JAX/Flax treinados no formato saved_model , que pode ser usado perfeitamente no TFX com treinamento genérico e avaliação de modelo. Para obter detalhes, verifique este exemplo .
scikit-aprender
Scikit-learn é uma biblioteca de aprendizado de máquina para a linguagem de programação Python. Temos um exemplo e2e com treinamento e avaliação customizados em TFX-Addons.