[논문 리뷰] EagerPy: Writing Code That Works Natively with PyTorch, TensorFlow, JAX, and NumPy

Jonas Rauber, Matthias Bethge외 1명|arXiv (Cornell University)|2020. 08. 10.

Computational Physics and Python Applications참고 문헌 13인용 수 4

한 줄 요약

EagerPy는 성능 오버헤드 없이 PyTorch, TensorFlow, JAX 및 NumPy와 함께 직접 작동하는 프레임워크에 종속되지 않는 딥러닝 코드를 작성할 수 있도록 해주는 파이썬 프레임워크입니다. 이는 투명한 디스패치를 통해 프레임워크 간 API를 통합하고, 메서드 체이닝을 지원하며, 포괄적인 타입 애너테이션을 제공하여 한 번 구현한 라이브러리가 여러 프레임워크에서 전체 성능과 타입 안정성을 유지하면서 사용 가능하게 합니다.

ABSTRACT

EagerPy is a Python framework that lets you write code that automatically works natively with PyTorch, TensorFlow, JAX, and NumPy. Library developers no longer need to choose between supporting just one of these frameworks or reimplementing the library for each framework and dealing with code duplication. Users of such libraries can more easily switch frameworks without being locked in by a specific 3rd party library. Beyond multi-framework support, EagerPy also brings comprehensive type annotations and consistent support for method chaining to any framework. The latest documentation is available online at https://eagerpy.jonasrauber.de and the code can be found on GitHub at https://github.com/jonasrauber/eagerpy.

연구 동기 및 목표

• 딥러닝 라이브러리 개발자가 PyTorch, TensorFlow, JAX 및 NumPy용 별도의 구현을 유지할 필요를 제거하기 위해.
• 사용자가 단일 라이브러리에 종속되지 않고 딥러닝 프레임워크를 자유롭게 전환할 수 있도록 하기 위해.
• 기본 프레임워크의 네이티브 성능을 그대로 유지하면서도 통합적이고 고성능이며 타입 안정적인 API를 제공하기 위해.
• 프레임워크에 종속되지 않는 코드를 프레임워크 전용 코드와 원활하게 통합하여 점진적인 도입을 가능하게 하기 위해.
• 딥러닝 프레임워크 간 코드 유지보수성 향상과 중복 제거를 위해.

제안 방법

• EagerPy는 중간 단계의 NumPy 변환 없이도, 기반 프레임워크의 네이티브 텐서 연산으로 직접 작업을 매핑하는 통합된 API를 도입합니다.
• CPU-GPU 메모리 이동 비용을 피하기 위해, 원본 프레임워크의 텐서에 직접 작업을 위임함으로써 네이티브 성능을 유지합니다.
• 모든 연산을 텐서 객체의 메서드로 노출시켜 메서드 체이닝을 지원함으로써, x.square().sum().sqrt().과 같은 직관적이고 읽기 쉬운 코드를 가능하게 합니다.
• MyPy를 사용하여 포괄적인 타입 애너테이션을 구현함으로써, 모든 함수 매개변수와 반환값에 대한 정적 타입 검사와 조기 버그 탐지 기능을 제공합니다.
• PyTorch, TensorFlow 2, JAX 및 NumPy를 포함한 여러 백엔드를 지원하기 위해 각각의 운영 체제에 맞는 프레임워크 전용 번역 레이어를 정의합니다.
• 엔드 투 엔드 테스트를 통해 일관된 동작과 API 정확성을 보장하는 테스트 기반 개발 방식을 사용하여, 모든 지원되는 프레임워크 간 일관성을 확보합니다.

실험 결과

연구 질문

• RQ1 PyTorch, TensorFlow, JAX 및 NumPy를 성능 저하 없이 직접 지원하는 단일 통합 API를 구축할 수 있는가?
• RQ2 서로 다른 문법적 및 의미적 관례를 가진 프레임워크 간에 메서드 체이닝과 일관된 API 설계를 어떻게 달성할 수 있는가?
• RQ3 포괄적인 타입 애너테이션이 딥러닝 라이브러리의 코드 안정성과 유지보수성 향상에 어느 정도 기여할 수 있는가?
• RQ4 기존 라이브러리에서 전체 재작성 없이도 프레임워크에 종속되지 않는 코드를 점진적으로 도입할 수 있는가?
• RQ5 PyTorch, TensorFlow 및 JAX 간의 자동 미분 API를 일관된 고수준 인터페이스로 추상화할 수 있는가?

주요 결과

• EagerPy는 네이티브 텐서에 직접 위임함으로써, 성능 오버헤드 없이 PyTorch, TensorFlow, JAX 및 NumPy에서 직접 실행되는 프레임워크에 종속되지 않는 코드를 작성할 수 있도록 합니다.
• 메서드 체이닝을 성공적으로 지원하여, 모든 지원되는 백엔드에서 x.square().sum().sqrt().과 같은 직관적이고 유창한 코드를 구현할 수 있습니다.
• MyPy를 사용하여 구현 및 검증된 포괄적인 타입 애너테이션은 코드 안정성 향상과 개발자 신뢰도 향상에 기여합니다.
• EagerPy는 Foolbox Native의 기초로 성공적으로 채택되어, 단일 코드베이스로 여러 프레임워크에서 네이티브 성능과 완전한 타입 안정성을 제공하게 되었습니다.
• EagerPy의 테스트 기반 개발 방식은 모든 지원되는 프레임워크 간 일관된 동작과 API 정확성을 보장하여, 버그 복귀 리스크를 줄였습니다.
• GUDHI와 같은 라이브러리들이 EagerPy를 채택하여, 코드 중복 없이도 PyTorch, TensorFlow 및 JAX에서 자동 미분 지원을 추가할 수 있게 되었습니다.