O TFX como plataforma é neutro em relação ao framework e pode ser usado com outros frameworks de ML, por exemplo, JAX, scikit-learn.
Para desenvolvedores de modelo, 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 de pipeline do TFX e a maioria dos módulos do TFX, por exemplo, orquestrador de pipeline, não dependem diretamente do TensorFlow, mas há alguns aspectos orientados ao TensorFlow, como formatos de dados. Com alguma consideração das necessidades de uma estrutura de modelagem específica, um pipeline do 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ê-lo como Arrow/RecordBatch na memória, que pode ser 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 de 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 de 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 de modelo ao servir. - O Trainer oferece suporte ao GenericTraining para que os usuários possam treinar seus modelos usando qualquer estrutura de ML.
- O avaliador , por padrão, oferece suporte apenas a
saved_model, mas os usuários podem fornecer uma UDF que gere previsões para avaliação do modelo.
Treinar um modelo em uma estrutura não baseada em Python exigirá o isolamento de um componente de treinamento personalizado em um contêiner do 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 pesquisa de aprendizado de máquina de alto desempenho. Flax é uma biblioteca de rede neural e ecossistema para JAX, projetado para 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-learn
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.