[논문 리뷰] TensorFlow Eager: A Multi-Stage, Python-Embedded DSL for Machine Learning

Akshay Agrawal, Akshay Naresh Modi외 9명|arXiv (Cornell University)|2019. 02. 27.

Parallel Computing and Optimization Techniques참고 문헌 26인용 수 52

한 줄 요약

TensorFlow Eager는 JIT 트레이서를 갖춘 TensorFlow에 대한 명령형 프런트엔드를 제공하며, Python 함수들을 데이터플로우 그래프로 스테이징할 수 있어 명령형 실행과 그래프 실행 간의 매끄러운 보간을 ML 워크로드에 가능하게 한다.

ABSTRACT

TensorFlow Eager is a multi-stage, Python-embedded domain-specific language for hardware-accelerated machine learning, suitable for both interactive research and production. TensorFlow, which TensorFlow Eager extends, requires users to represent computations as dataflow graphs; this permits compiler optimizations and simplifies deployment but hinders rapid prototyping and run-time dynamism. TensorFlow Eager eliminates these usability costs without sacrificing the benefits furnished by graphs: It provides an imperative front-end to TensorFlow that executes operations immediately and a JIT tracer that translates Python functions composed of TensorFlow operations into executable dataflow graphs. TensorFlow Eager thus offers a multi-stage programming model that makes it easy to interpolate between imperative and staged execution in a single package.

연구 동기 및 목표

• ML 워크플로우를 위한 명령형 사용성과 그래프 기반 최적화를 결합하는 DSL의 필요성을 동기 부여한다.
• 기본적으로 명령형 실행을 지원하고 선택적 그래프 스테이징을 제공하는 Python 내장 다단계 프로그래밍 모델을 제안한다.
• Python 함수를 실행 가능한 데이터플로우 그래프로 변환하는 옵트인 트레이싱 기반 JIT의 구현을 설명한다.
• 실행 모드 간 원활한 전환을 가능하게 하도록 상태, 디바이스, 분산이 어떻게 표현되는지 설명한다.
• 대표 ML 모델에서 명령형 실행과 스테이징 실행 간의 성능 트레이드오프를 평가한다.

제안 방법

• TensorFlow 위에 계층화된 Python 내장 DSL을 도입하여 기본적으로 명령형으로 실행할 수 있고 스테이징 시 Python 함수를 트레이스하여 데이터플로우 그래프를 구축한다.
• 입력/출력에 이름이 있는 그래프 함수를 생성하기 위해 Python 함수를 트레이스하는 함수 데코레이터와 실행을 위한 C++ 데이터플로우 런타임을 제공한다.
• 중첩된 테이프를 사용한 고阶 도함수용 추적 기반 역전파 자동 미분에 대해 설명한다.
• 직렬화 및 복원을 위한 변수 및 그래프 기반 상태 매칭을 통한 상태 관리에 대해 설명한다.
• 자동 디바이스 배치, 디바이스 간 데이터 전송, TPU/XLA 기반의 스테이징 경로를 포함한 디바이스 및 분산 지원에 대해 논의한다.
• 그래프 내에 명령형 코드를 포함하고 추적 중 상태 생성을 관리하기 위한 이스케이핑 및 통합 메커니즘(py_func, tf.init_scope)을 개요한다.

실험 결과

연구 질문

• RQ1 명령형 친화적 프런트엔드를 TensorFlow에 설계하여 그래프 기반 실행과의 보완 없이 데이터플로우 그래프의 이점을 유지할 수 있는가?
• RQ2 명령형 실행과 스테이징된 그래프 구성이라는 다단계 프로그래밍 모델의 설계, 구현, 성능 트레이드오프는 무엇인가?
• RQ3 상태, 디바이스 배치, 분산이 두 실행 모드에서 일관되게 관리될 수 있는가?
• RQ4 트레이싱 기반 자동 미분이 부분적으로 스테이징된 계산 및 Python의 제어 흐름과 어떻게 상호작용하는가?

주요 결과

• 명령형 TensorFlow Eager는 명령형 실행을 사용할 때 단일 GPU에서 ResNet-50의 그래프 성능에 상응할 수 있다.
• 스테이징된 TensorFlow Eager는 TPU 가속 학습에서 ResNet-50을 명령형 TensorFlow Eager가 TPU에서 달성할 수 있는 것보다 훨씬 빠르게 수행한다.
• 작은 연산을 갖는 모델의 경우 스테이징이 상당한 속도향상을 가져오며 코드 변경은 최소화된다.
• 프레임워크는 실행 모드에 무관한 일관된 단일 API 표면을 제공하고 도구 및 배포를 위해 TensorFlow 생태계를 활용한다.
• 구현은 경량적이며(Python 수준의 스테이징 약 2000줄, 핵심은 Python과 C에서 차별화) 크로스 플랫폼이다.