TFX 머신러닝 파이프라인 사용하기

새 데이터로 학습·평가를 반복하는 지속 학습 파이프라인을 TFX로 만들면서 사내 Dataset Registry와 Model Registry, Apache Beam을 연결한 커스텀 컴포넌트 구현기다. ComponentSpec, Interface, Executor, Artifact와 ML Metadata의 역할부터 최신 데이터 감지, TFRecord 변환, 클라우드 학습, Inferentia compile과 모델 등록까지 이어진다.

핵심 포인트

TFX 컴포넌트는 입출력·파라미터를 선언하는 ComponentSpec, 실행 로직인 Executor, 둘을 묶는 Interface로 구성된다.
컴포넌트 간 데이터는 Channel의 Artifact로 이동하며 메타데이터는 MLMD, 실제 파일은 로컬이나 GCS 같은 storage에 저장된다.
사내 Dataset Registry를 읽는 ExampleGen을 BaseBeamComponent와 BaseExampleGenExecutor로 구현해 형식이 다른 데이터를 TFRecord로 바꾼다.
Dataset Registry가 파일 시스템이 아니므로 InputProcessor를 확장해 최신 span과 version을 찾아 새 데이터가 들어올 때 파이프라인을 실행한다.
학습 결과는 SavedModel로 만들고 커스텀 컴포넌트가 필요하면 Inferentia Neuron compile까지 수행한 뒤 Model Registry에 등록한다.

상세 정리

도입 배경: 새 데이터를 받을 때마다 학습과 평가를 사람이 실행하는 비효율을 줄이고 지속 학습 과정을 재현 가능한 파이프라인으로 만들려 했다.
TFX 범위: 데이터 가공, 학습, 검증뿐 아니라 실제 환경과 유사한 요청 검증, warmup 데이터 생성, 서빙까지 지원한다. 팀은 필요한 일부만 사내 시스템에 맞게 고쳤다.
Spec 역할: INPUTS와 OUTPUTS는 ChannelParameter, 설정값은 ExecutionParameter로 선언한다. 지원 기본형이나 protobuf만 받을 수 있다.
직렬화 함정: protobuf와 JSON 설정은 내부적으로 문자열로 전달되므로 Executor에서 다시 역직렬화해야 한다.
Interface 역할: SPEC_CLASS와 EXECUTOR_SPEC을 연결하고 출력 Artifact Channel을 만든다. 초기화 단계에서는 실제 데이터를 읽거나 전처리할 수 없다.
Executor 입력: `input_dict`, `output_dict`, `exec_properties`가 각각 입력 Artifact, 출력 Artifact, 실행 파라미터를 제공한다.
저장소 추상화: Artifact의 `uri`가 로컬 경로뿐 아니라 GCS일 수 있으므로 Python `open()` 대신 GFile이나 TFX io_utils를 사용한다.
Beam 통합: Interface와 Executor가 각각 BaseBeamComponent와 BaseBeamExecutor를 상속하고 `_make_beam_pipeline()`에서 파이프라인을 작성한다.
실행 인자: 컴포넌트에 `with_beam_pipeline_args()`를 붙여 Dataflow 등 Beam runner 옵션을 개별 전달할 수 있다.
Dataset 변환: Registry는 CSV와 JSONL 등 여러 형식을 받으므로 사용자 transform 함수를 적용해 직렬화하고 `tf.train.Example`로 변환한다.
ExampleGen 확장: BaseExampleGenExecutor의 source-to-example PTransform을 구현해 Registry 읽기, 파싱, TFRecord 변환을 연결했다.
새 버전 탐지: InputProcessor의 `get_latest_span()`과 `get_latest_version()`이 Registry 목록에서 최신 버전을 찾는다. 패턴에 `{SPAN}_{VERSION}`을 넣어 같은 파이프라인을 재실행한다.
스키마·변환: StatisticsGen과 SchemaGen이 TFRecord 스키마를 추론하고 필요하면 ImportSchemaGen으로 고정한다. Transform은 주로 tokenizer 처리에 쓴다.
학습 분기: 로컬은 기본 Trainer, 클라우드는 Google Cloud AI Platform Trainer를 사용한다. 서빙용과 평가용 Signature를 분리한 SavedModel을 낸다.
모델 등록: 커스텀 Executor가 SavedModel을 읽고 필요하면 `tensorflow.neuron.trace`로 변환한 뒤 버전 정보와 함께 Model Registry에 올린다.