딥러닝 프레임워크

'텐서플로우 vs 파이토치 vs JAX' 비교

오늘날 딥러닝 연구 및 개발을 주도하는 3가지 프레임워크가 있다. 각각 사용 편의성, 기능 및 성숙도, 엄청난 확장성으로 유명하다. 어떤 프레임워크를 사용해야 할까? 
시리, 알렉사, 실시간 번역 앱, 로봇, 자율주행차 등에 들어가는 ‘딥러닝’은 크고 작은 방식으로 삶을 매일 변화시키고 있다. 대부분의 딥러닝 애플리케이션은 텐서플로우, 파이토치, JAX 중 하나를 이용해 작성된다. 그렇다면 개발자 입장에서 기술을 만들 때 어떤 딥러닝 프레임워크를 선택해야 할까?

https://jcoding7.tistory.com/entry/%EB%94%A5%EB%9F%AC%EB%8B%9D

딥러닝

https://jcoding7.tistory.com/entry/%EA%B0%9C%EB%B0%9C%EC%9E%90-%EC%9A%A9%EC%96%B4-%ED%94%84%EB%A0%88%EC%9E%84%EC%9B%8C%ED%81%AC

개발자 용어 - 프레임워크

‘텐서플로우’의 장점

1970년대와 1980년대에는 ‘IBM 제품을 샀다고 해고된 사람은 아무도 없(Nobody ever got fired forbuying an IBM)’라는 말은 진리로 통했다. 이제 이 표현을 ‘딥러닝에 텐서플로우를 사용했다고 해고된 사람은 아무도 없다’라고 바꿔 말할 수 있겠다. 하지만 IBM이 1990년대에 들어서면서 도태한것 처럼 공개된 지 7년이 지난 텐서플로우는 지금에도 여전히 경쟁력이 있을까?

확실히 그렇다. 텐서플로우는 그동안 가만히 있지 않았다. 텐서플로우 1.x는 파이썬과 매우 다른 방식으로 정적 그래프를 작성하는 게 전부였지만, 텐서플로우 2.x는 ‘즉시 실행’ 모드를 사용한 모델 빌드가 가능해지면서 좀 더 파이토치 같은 느낌이 났다.

아울러 하이 레벨(high level)에서 텐서플로우는 더 쉬운 개발을 위해 케라스(Keras)를 제공하고, 로우 레벨에서는 속도를 위해 컴파일러를 최적화하는 XLA(Accelerated Linear Algebra)를 제공한다. XLA는 GPU 성능을 향상시키며, 대규모 모델 학습에 뛰어난 성능을 제공하는 구글의 TPU(Tensor Processing Units)를 활용하는 방법이기도 하다.

성숙한 플랫폼에서 잘 정의되고 반복 가능한 방식으로 모델을 제공해야 한다면 ‘텐서플로우 서빙(TensorFlow Serving)’을 이용하면 좋다. 웹, 스마트폰 등 저전력 컴퓨팅 또는 IoT 기기처럼 리소스가 제한된 장치용으로 모델 배포 대상을 재설정해야 할 경우, 텐서플로우.js(TensorFlow.js)나 텐서플로우 라이트(TensorFlow Lite)를 이용하면 된다. 구글이 텐서플로우를 사용해 프로덕션 수준의 제품을 100%를 실행하고 있다는 사실을 감안하면, 텐서플로우는 규모가 큰 곳에서도 활용하기 좋다는 것을 예측할 수 있다.

하지만 이 프로젝트를 둘러싼 지원은 점점 부족해지고 있다. 텐서플로우 1.x에서 2.x으로의 업그레이드는 한마디로 혹독했다. 몇몇 기업들은 새 버전에서 제대로 작동하도록 코드를 업데이트하는 데 필요한 수고로움을 검토한 결과, 코드를 파이토치로 이식하는 쪽을 택하기도 했다. 텐서플로우는 연구계에서도 (몇 년 전부터 파이토치가 제공하는 유연성이 선호되기 시작하면서) 관심을 잃었고, 결과적으로 텐서플로우를 사용하는 연구 논문도 줄었다.

케라스 사건도 도움이 되지 않았다. 케라스는 2년 전 텐서플로우 릴리스에 통합됐지만 최근에 다시 한번 독자적인 릴리스 일정이 있는 별도의 라이브러리로 분리됐다. 물론 케라스의 분리가 개발자의 삶에 영향을 미치는것은 아니지만 무시할만한 변화는 아니다. 

그래도 어쨌든 텐서플로우는 신뢰할 수 있는 프레임워크이자 딥러닝을 지원하는 광범위한 생태계의 호스트다. 텐서플로우에서는 모든 규모에서 작동하는 애플리케이션과 모델을 구축할 수 있다. 하지만 요즘에는 텐서플로우가 가장 첫 번째 선택지는 아닐 수도 있다.

‘파이토치’의 장점

한때 텐서플로우의 뒤를 바짝 뒤쫓는 건방진 신예였던 파이토치는 현재 딥러닝 세계의 핵심 세력이 됐다. 아마도 대부분 연구용으로 쓰이지만 프로덕션 애플리케이션용으로도 점점 더 많이 사용되고 있다. 즉시 실행 모드가 파이토치뿐만 아니라 텐서플로우에서도 기본 개발 방법이 된 가운데, 파이토치의 자동 미분(Automatic Differentiation/Autograd)이 제공하는 더 파이썬스러운(Pythonic) 접근 방식이 정적 그래프와의 전쟁에서 승리한 것으로 보인다.

텐서플로우와 달리 파이토치는 버전 0.4에서 변수 API(Variable API) 지원을 중단한 이후 심각한 핵심 코드 파열을 전혀 겪지 않았다(예전에는 텐서와 자동 미분을 사용하려면 변수가 필요했는데 지금은 모든 것이 텐서다). 그렇다고 실수가 없었던 건 아니다. 예를 들어 파이토치를 사용해 여러 GPU에서 학습시켜 왔다면 데이터패러랠(DataParellel)과 최신 디스트리뷰티드데이터패러랠(DistributedDataParallel) 간의 차이를 발견했을 것이다. 거의 항상 후자를 사용해야 하지만 전자가 실제로 사용 중단된 것은 아니다. 

파이토치는 XLA/TPU 지원 측면에서 텐서플로우 및 JAX보다 뒤처졌었지만 2022년 현재는 상황이 크게 개선됐다. 이제 파이토치는 구형 스타일의 TPU 노드 지원은 물론 TPU VM 액세스 지원도 제공한다. 이와 함께 코드 변경 없이 CPU, GPU 또는 TPU에서 코드를 실행할 수 있는 간편한 명령줄 구현도 지원한다. 

아울러 파이토치에서 작성해야 할 때가 많은 보일러플레이트 코드를 처리하고 싶지 않다면 파이토치 라이트닝(PyTorch Lightning)과 같은 하이 레벨 추가 기능을 사용하면 된다. 파이토치 라이트닝을 활용하면 학습 루프를 재작성하는 대신 실제 작업에 집중할 수 있다. 단, 파이토치 모바일(PyTorch Mobile)에서 작업이 계속되지만 텐서플로우 라이트보다는 미숙하다는 게 단점이다. 

프로덕션 측면에 봤을 때 파이토치는 이제 프레임워크에 구애받지 않는 플랫폼(예: 큐브플로우(Kubeflow) 등)과 통합된다. 또한 토치서브(TorchServe) 프로젝트를 사용하면 확장, 지표, 일괄 추론 등의 배포 세부 사항을 처리할 수 있다. 파치토치 개발자가 직접 관리하는 작은 패키지로 ML옵스(MLOps)의 모든 장점을 누릴 수 있는 셈이다. 

파이토치는 확장될까? 메타는 수년간 파이토치를 프로덕션 환경에서 실행하고 있기 때문에 파이토치가 대규모 워크로드를 처리할 수 없다는 말은 거짓말이다. 그래도 GPU 또는 TPU가 많이 필요한 대규모 학습용으로는 파이토치가 JAX만큼 친화적이지는 않다고 할 수 있다.

마지막으로 금기시되는 주제가 있다. 지난 몇 년간 파이토치의 인기는 허깅 페이스(Hugging Face)의 트랜스포머(Transformers) 라이브러리의 성공과 밀접하게 관련돼 있다. 물론 트랜스포머는 (이제) 텐서플로우와 JAX도 지원하지만 원래는 파이토치 프로젝트로 시작했고, 여전히 이 프레임워크와 긴밀하게 결합돼 있다.

트랜스포머 아키텍처의 부상, 파이토치의 유연성, 허깅 페이스의 모델 허브를 통해 공개 후 불과 며칠 내지 또는 몇 시간 만에 새로운 모델을 끌어들일 수 있는 기능 등을 감안하면 파이토치가 오늘날 어디서나 인기를 얻고 있는지 쉽게 알 수 있다.

‘JAX’의 장점

텐서플로우에 별로 관심 없는 사용자를 위해 구글은 일종의 대안을 마련해 뒀다. JAX는 구글에서 개발하고 유지관리하며 사용하는 딥러닝 프레임워크이지만 공식적으로 구글 제품은 아니다.

하지만 구글/딥마인드에서 지난 약 1년 동안 나온 논문과 릴리스를 살펴보면 구글 많은 연구가 JAX로 옮겨갔다는 사실을 눈치챌 수밖에 없다. 그러니까 JAX는 ‘공식’ 구글 제품은 아니지만 구글 연구진이 한계를 뛰어넘기 위해 사용하고 있는 제품이다.

그렇다면 JAX는 정확히 무엇인가? 이렇게 생각하면 간단하다. 요술 지팡이를 휘두르면 마법처럼 파이썬 함수를 벡터화할 수 있고, 해당 함수 상의 모든 파생 계산을 처리할 수 있는 GPU/TPU 가속 버전의 넘파이(NumPy)를 상상해 보라. 또한 코드를 가져와 XLA 컴파일러에 맞게 최적화하는 JIT(Just-In-Time) 구성요소가 있어 텐서플로우와 파이토치에 비해 성능이 크게 향상됐다. 실제로 최적화 없이 JAX에서 재구현하는 것만으로 일부 코드 실행 속도가 4~5배 높아지는 것을 확인한 바 있다.

JAX가 넘파이 수준에서 작동한다는 점을 감안하면 JAX 코드는 텐서플로우/케라스는 물론 파이토치보다 훨씬 낮은 수준에서 작성된다. 다행히도 이 프로젝트를 중심으로 한 작지만 성장 중인 생태계가 있다. 신경망 라이브러리를 원하는가? 그렇다면 구글의 플랙스(Flax) 와 딥마인드(역시 구글)의 하이쿠(Haiku)가 있다. 이 밖에도 온갖 최적화 니즈를 충족할 옵택스(Optax)와 이미지 처리를 위한 PIX 등이 있다. 

일단 플랙스 등으로 작업하게 되면 신경망 구축은 비교적 이해하기 쉬워진다. 단, 몇몇 어려운 부분이 여전히 있다는 점만 유의하면 된다. 예를 들어 베테랑들은 JAX의 난수 처리 방식이 다른 여러 프레임워크와 다르다는 점을 많이 언급한다. 모든 것을 JAX로 바꿔야 할까? 학습에 막대한 리소스가 필요한 대규모 모델 연구에 깊이 관여 중이라면 그렇다. 결정 학습 및 수천 개의 TPU 포드가 필요한 상황이라면 JAX가 이뤄낸 발전만으로도 바꿀 가치가 충분하다.

텐서플로우 vs. 파이토치 vs. JAX

결론은 무엇인가? 어떤 딥러닝 프레임워크를 사용해야 할까? 아쉽게도 정답은 없다. 이는 해결하려는 문제의 종류, 배포 할 모델의 크기, 타깃 컴퓨팅 플랫폼 등에 따라 달라질 수 있다. 

하지만 텍스트 및 이미지 영역에서 작업 중이고, 이러한 모델을 프로덕션에 배포할 목적으로 중간 규모의 연구를 하고 있다면 아마도 파이토치가 최선의 선택이라고 해도 무방하다. 요즘 그 분야에서는 파이토치가 최고다. 단, 저사양 기기에서 성능을 뽑아내야 한다면 든든한 텐서플로우 라이트 패키지가 있는 텐서플로우를 추천한다. 반대로, 연구 목적으로 매개변수가 수백억 내지 수천억 개 이상인 모델을 학습시키고 있다면 JAX를 사용해 보는 것도 좋다.