[논문 리뷰] Learning to Optimize Tensor Programs

Tianqi Chen, Lianmin Zheng외 6명|arXiv (Cornell University)|2018. 05. 21.

Parallel Computing and Optimization Techniques참고 문헌 36인용 수 124

한 줄 요약

논문은 딥 러닝 워크로드를 위한 텐서 프로그램 구현을 자동으로 검색하고 최적화하기 위해 도메인 특화 비용 모델을 구축하는 학습 기반 프레임워크 AutoTVM을 제시합니다. 전이 학습을 통해 CPU, GPU, 모바일 하드웨어 전반에서 경쟁력 있는 성능을 달성합니다.

ABSTRACT

We introduce a learning-based framework to optimize tensor programs for deep learning workloads. Efficient implementations of tensor operators, such as matrix multiplication and high dimensional convolution, are key enablers of effective deep learning systems. However, existing systems rely on manually optimized libraries such as cuDNN where only a narrow range of server class GPUs are well-supported. The reliance on hardware-specific operator libraries limits the applicability of high-level graph optimizations and incurs significant engineering costs when deploying to new hardware targets. We use learning to remove this engineering burden. We learn domain-specific statistical cost models to guide the search of tensor operator implementations over billions of possible program variants. We further accelerate the search by effective model transfer across workloads. Experimental results show that our framework delivers performance competitive with state-of-the-art hand-tuned libraries for low-power CPU, mobile GPU, and server-class GPU.

연구 동기 및 목표

• 텐서 연산자를 수동으로 튜닝된 라이브러리 너머 자동 최적화의 동기를 제공한다.
• 하드웨어 런타임을 최소화하기 위한 거대한 텐서 프로그램 스케줄 공간 탐색 문제를 형식화한다.
• 수십억 개의 프로그램 변형을 탐색하기 위한 학습 기반 비용 모델을 개발한다.
• 다른 워크로드와 하드웨어 대상 간의 지식을 재활용하기 위한 전이 학습을 가능하게 한다.

제안 방법

• 주어진 지수 표현 e 와 코드 생성기 g에 대해 비용 f(x)가 알려져 있지 않지만 측정 가능한 스케줄 검색 공간 Se 를 정의한다.
• 저수준 AST로부터의 도메인 특성들을 사용하는 그래디언트 부스팅 트리(GBT)와 AST의 TreeGRU 임베딩의 두 가지 비용 모델을 개발한다.
• 런타임을 예측하고 선택을 유도하기 위해 회귀 혹은 랭크 기반 목표로 모델을 학습한다.
• 시뮬레이션된 어닐링과 다양성 인식을 통한 선택으로 후보 스케줄을 생성하고 하드웨어에서 실제 성능을 측정하는 탐색 루프를 사용한다.
• 전역 불변 표현 모델과 로컬 도메인 내 모델을 결합한 전이 학습을 도입하여 새로운 워크로드에 대한 최적화를 가속화한다.
• AutoTVM(TVM 기반)이 외부 연산 라이브러리 없이 백엔드 전반에 걸쳐 엔드-투-엔드 이득을 달성한다고 보고한다.

실험 결과

연구 질문

• RQ1 학습 기반 비용 모델이 수십억 규모의 스케줄 공간을 실 hardware에서 빠른 텐서 프로그램 구현으로 효과적으로 탐색할 수 있는가?
• RQ2 전이 학습이 서로 다른 워크로드와 하드웨어 대상 간의 최적화를 더 빠르게 가능하게 하는가?
• RQ3 이 도메인에서 빠른 스케줄을 찾는 데 어떤 목적(회귀 대 랭크)이 가장 잘 작동하는가?
• RQ4 다양성 인식 탐색과 불확실성 추정이 이 설정에서 유익한가?

주요 결과

• 통계적 비용 모델(GBT 및 TreeGRU)은 실제 하드웨어 평가를 더 적게 수행하면서도 더 빠른 텐서 프로그램을 찾아내는 데 블랙박스 기준선보다 우수하다.
• 랭크 기반 목표는 일반적으로 후보 프로그램의 속도 순서를 매기는 데 회귀 목표보다 동등하거나 더 잘 작동한다.
• 다양성 인식 탐색은 거의 영향이 없거나 미미했으며, 특정 워크로드에서 일부 이익을 보였다.
• 전이 학습은 GPU 및 모바일 장치 전반에 걸쳐 이전에 본 워크로드의 지식을 재사용하여 2배에서 10배의 속도 향상을 제공한다.
• NVIDIA TITAN X, ARM Cortex-A53, ARM Mali 백엔드 전반에서 AutoTVM의 엔드-투-엔드 추론 향상은 1.2배에서 3.8배 범위이다.