[살아 움직이는 머신러닝 파이프라인 설계] 2. TFX - 텐서플로 익스텐디드

TensorFlow eXtended (TFX)

파이프라인을 직접 작성하면 가장 문제가 많이 발생하는 부분이 글루 코드다. 글루 코드란 파이프라인 각 단계를 연결하는 코드를 말한다. 머신러닝 프로젝트에서 데이터 수집부터 모델 배포까지 각 단계를 거쳐야 하는데, 어느 단계에서 코드를 수정한다면 이를 파이프라인에 적용할 때 정상적으로 작동되는지 확인해야 한다. 만약에 각 단계마다 의존성이 심하면 버그를 해결하기 위해 다른 단계도 손봐야 할 수 있다. 이러한 일이 계속 발생한다면 머신러닝 파이프라인은 유지보수되기 어렵다.
구글에서도 이러한 문제를 계속 겪어왔다. 그래서 이러한 일을 해결하기 위해 Tensorflow Extended를 만들었다. TFX로 파이프라인의 정의를 단순화하고 기본적으로 필요한 코드의 양을 최소화하여 글루 코드로 인해 발생하는 문제를 방지하고 유지보수도 편리하게 만들었다.

TFX를 이용한 머신러닝 아키텍쳐

TFX을 사용한 머신러닝 아키텍쳐의 모습

프로젝트에 따라 머신러닝 파이프라인은 단계들이 반드시 한 단계씩 일어나지 않을 수 있다. 때로는 각 단계의 처리 순서가 유향 비순환 그래프 형태를 띨 수 있다. 또한 커다란 데이터셋을 처리해야하는 경우도 발생할 수 있다. 그래서 필요에 따라 각 단계를 병렬 처리할 수 있어야 한다. 그렇기 때문에 TFX에서 머신러닝 아키텍쳐를 구성할 때 파이프라인 오케스트레이션과 메타데이터 스토어를 두게 된다.

파이프라인 오케스트레이션 (Pipeline Orchestration)

파이프라인 오케스트레이션이란 데이터를 처리하는 각 과정들이 병렬적으로 원할하게 처리하도록 파이프라인을 관리하는 도구를 말한다. 파이프라인 오케스트레이션을 통해 머신러닝 파이프라인의 각 과정을 병렬적으로 실행하게 된다. TFX는 Apache Beam을 사용해 구현했기 때문에 Apache Airflow, Kubeflow Pipeline 등, 원하는 병렬 처리 러너를 사용할 수 있다.

메타데이터 스토어

TFX에서는 머신러닝 파이프라인의 각 과정에서 들어오거나 처리되는 모든 것들을 아티펙트(Artifacts) 라 부른다. 원시 데이터나 전처리된 데이터, 심지어 학습된 모델까지 모두 아티펙트라고 칭한다. 하지만 TFX에서는 각 과정이 아티펙트를 직접 주고받지 않는다. 모든 아티펙트를 한 데이터베이스에 저장하고 각 아티펙트에 대한 참조를 주고받는다. 여기서 아티펙트에 대한 참조를 담은 정보를 메타데이터(Metadata) 라 부르고 메타데이터 스토어를 두어 각 과정에서 주고받는 메타데이터를 관리하고 감시한다.

TFX가 제공하는 요소

TFX가 포함하는 TFX 컴포넌트와 라이브러리

TFX는 머신러닝 파이프라인을 구성하기 위해 크게 두가지를 제공한다. 하나는 각 단계에서 실질적으로 작업을 처리하는 함수나 클래스가 모인 라이브러리이고 나머지는 이러한 라이브러리를 사용해 파이프라인의 한 단계를 대표하는 TFX 컴포넌트다.
라이브러리는 익히 우리가 잘 알고있는 그 의미 그대로다. 본래 구글에서는 처음에 TFX를 그대로 제공하지 않고 각 파이프라인 과정을 구성할 때 사용했던 라이브러리를 따로 제공했다. 텐서플로우로 학습한 모델을 쉽게 배포할 수 있도록 도와주는 Tensorflow Serving부터 데이터 검증에 사용되는 Tensorflow Data Validation, 데이터 전처리를 위한 Tensorflow Tranform, 모델의 성능을 비교하기 위한 Tensorflow Model Analysis까지 모두 TFX로 흡수되 하나의 플랫폼을 구성하고 있다.
TFX 컴포넌트는 이러한 라이브러리를 통해 각 단계에서 필요한 처리를 담당한다. 보통 머신러닝 파이프라인을 구현한다면 라이브러리를 통해 직접 구현하기 보다 TFX 컴포넌트를 통해 쉽고 편하게 구현할 것이다. 무엇보다 TFX 컴포넌트는 파이프라인 오케스트레이션에게 자동으로 필요한 작업을 전달하기 때문에 파이프라인 외부에서 테스트를 하지 않는 이상 라이브러리 만을 사용하는 일은 거의 없다.

TFX Component로 구현된 머신러닝 파이프라인과 데이터 흐름, 출처: https://www.tensorflow.org/tfx?hl=ko

TFX 컴포넌트 구조

TFX 컴포넌트의 구조, 출처: https://www.tensorflow.org/tfx/guide/understanding_custom_components

TFX 컴포넌트는 크게 3가지로 분리된다. 메타데이터를 받는 드라이버(Driver), 메타데이터를 통해 가져온 데이터를 처리하는 실행자(Executor), 처리된 데이터를 가리키는 메타데이터를 메타데이터 스토어에 반환하는 배포자(Publisher)로 구성된다. 여기서 주목할 점은 위에서도 설명했지만 TFX 컴포넌트는 데이터를 직접 교환하지 않고 메타데이터를 통해 접근한다. 아무래도 직접 각 프로세스가 데이터를 주고받는 것보다 한곳에 저장하고 참조를 주고받는 쪽이 병렬처리로 구현하기 쉽기 때문일 것이다.

대화형 파이프라인

Colab에서 TFX 컴포넌트를 실행하는 모습, 출처: https://blog.tensorflow.org/2019/11/introducing-tfx-interactive-notebook.html

TFX가 제공하는 또 하나의 기능이 대화형 파이프라인이다. 이는 Jupyter같은 파이썬 노트북에서 쉽게 TFX 컴포넌트를 실행하고 결과를 바로 확인할 수 있는 기능이다. 머신러닝 파이프라인이 꽤 큰 프로젝트인 만큼 디버깅과 실험이 쉽도록 만들기 위해 이러한 기능을 제공하는 것 같다.