파이토치로 구현된 딥러닝 코드에 익숙해진 나는 최근에 다른 라이브러리도 적용해보고 싶다는 생각이 들었다. 파이토치의 유연함과 편리함 덕분에 많은 실습에서 활용해 왔지만, 새로운 프레임워크를 통해 딥러닝의 다양한 접근 방식을 경험해 보는 것도 필요하다고 느꼈다.

그러던 중 최근에 JAX와 Flax를 활용한 딥러닝 프레임워크에 대해 다룬 이 책을 접하게 되었고, 바로 읽어보게 되었다.

이 책은 모두의 연구소 JAX/Flax LAB이

JAX를 어떻게 실용적으로 활용할 수 있는지에 중점을 두고 집필한 책이다.

그 목표에 맞게 이론과 실제 예제가 적절하게 섞여있다.

개발 서적치고 짧은 분량에도 불구하고 간결하면서 쉽게 이해 가능한 설명과 예제 코드에 대한 설명으로 가득 채워져 있어서 빠르고 수월하게 JAX를 익히기에 좋은 책이었다.

전에 텐서플로와 파이토치로 구현해 본 적이 있던 CNN, ResNet, DCGAN 같은 친숙한 딥러닝 모델을 JAX로 구현해 보면서 해당 라이브러리를 더 쉽게 이해할 수 있었다.

Numpy와 JAX를 비교해서 설명하는 부분과 참고 내용으로 텐서플로, 파이토치에서는 어떻게 구현하는지 짧게 비교 설명한 부분도 마음에 들었다.

이 책은 파이썬에 익숙하며 머신러닝/딥러닝에 대한 사전 지식이 있고 파이토치 보다 확장성이 뛰어난 라이브러리를 프로젝트나 연구에 적용할 필요성을 느낀 사람들에게 추천한다.

JAX는 구글 딥마인드에서 사용하고 있는 고성능 수치계산을 지원하는 딥러닝 프레임워크이다. 뒤에서 또 언급하겠지만 JAX는 PyTorch나 TensorFlow와 다르게 함수형 프로그래밍 방식을 사용하기 때문에 이에 대한 이해가 좀 필요하다. 그리고, Flax는 구글 브레인 팀에서 JAX팀과 협업을 하면서 JAX를 조금 더 쉽게 사용할 수 있도록 하기 위해 만들어진 오픈소스로서, Flexibility+JAX의 의미로 만들어진 합성어이다. 그래서 요새 JAX 기반으로 딥러닝 모델을 만들 때에는 대부분 Flax를 함께 사용한다고 한다아래는 OpenAI에서 음성인식 및 음성 번역을 위해 PyTorch로 만든 오픈소스인 Whisper(첫번째 컬럼)와 허깅 페이스 연구원이 이를 JAX로 변환해서 만든 모델인 Whisper-JAX와의 성능 비교표이다. A100 GPU(난 아직 써본 적이 없는 1500만원 짜리 GPU..ㅠ) 환경에서 음성 데이터의 길이별 추론 성능이 거의 1/10 수준으로 Whispter-JAX의 처리 속도가 빠른 것을 확인할 수 있다. JAX의 성능이 PyTorch보다 상대적으로 빠르다는 걸 단적으로 보여주는 표라고 할 수 있겠다. 

벤치마크 출처 : https://github.com/sanchit-gandhi/whisper-jax

JAX 활용 기본 예제 코드

3장은 Flax에 대한 설명을 담은 chapter이다. 앞서도 소개했듯이 Flax는 JAX를 사용하는 딥러닝 연구자들을 위해 개발된 고성능 신경망 라이브러리이자 에코 시스템이다. 최초 구글 리서치 브레인 팀에서 JAX팀과 협력하며 개발하였고 현재는 오픈소스 커뮤니티와 공동으로 개발하고 있다. TensorFlow 딥러닝 프레임워크의 쉽고 간결한 사용을 위해 Keras API가 나온 것처럼, JAX 프레임워크의 사용을 돕기 위해 나온 것이 Flax라고 보면 된다. 

2~3장에서 JAX/Flax의 기본적인 사용법을 다루었다면, 4장에서는 이러한 기초 지식을 바탕으로 JAX/Flax 라이브러리를 통해 딥러닝 모델 구축과 미세조정(파인튜닝) 과정을 수행하는 내용에 대해 다루고 있다. JAX/Flax에 대한 가장 종합적인 활용 부분을 담고 있는 본 도서의 핵심 chapter라고 할 수 있겠다.

TRC 프로그램 신청 사이트 (https://sites.research.google/trc/about/)

[책리뷰] JAX/Flax로 딥러닝 레벨업

전반적인 소감 및 마음에 드는 부분

흔히들 딥러닝 프레임워크 하면 Tensorflow나 PyTorch를 떠올린다. 인공지능 개발 경쟁이 한창인 요즘 JAX는 구글 딥마인드에서 개발, 유지관리하고 사용하는 또 다른 딥러닝 프레임워크다. 텐서플로우나 파이토치를 알려주는 책은 많지만 JAX 를 알려주는 책이 24년 9월 기준 이 책이 유일하다. 이 책에서도 쓸데없는 파이썬 기초 프로그래밍을 알려주는 페이지 없이 JAX/Flax 에 대해서만 다룬다. 책의 머릿말에도 나오듯이 JAX/Flax LAB은 '모두의 연구소'에서 활동하는 AI 리서처와 현업자로 구성되어 있는데, 해당 구성원들이 집필한 책이다. 아무리 공식 문서를 보고 공부하는 게 좋다고 하더라도 이렇게 번역하고 한국어로 책이 나온 점은 대단히 감사하다.

대상 독자 및 책 난이도

JAX에 관심 있는 모든 개발자와 학생들을 대상으로 한다. 특히 파이썬을 통한 고성능 계산과 머신러닝에 익숙하거나 관심이 있는 분들도 대상 독자다. 이 책을 보기 위해 적어도 파이썬 프로그래밍 기초 지식이 있어야 하며, 선형대수 기본 개념과 머신러닝/딥러닝 기초 이론이 필요하다. 여기에 함수형 프로그래밍이 대략 어떤건지 간단하게나마 알고 있으면 더욱 좋다. 개인적으로는 텐서플로우와 파이토치로 딥러닝을 공부해본 개발자를 대상으로 한 느낌이다.

텐서플로우나 파이토치로 딥러닝 공부를 해본 경험이 있다면 무리없이 따라할 수 있다. 파이썬 라이브러리 import 할 때, JAX로 대체하는 것이 대부분이라고 보면 될 것이다. 그러나 처음부터 JAX로 시작한다고 하면 어려울 수 있다.

다루는 내용과 범위

JAX/Flax 프레임워크에 대해서만 다룬다. 책의 목차는 다음과 같다.

1. JAX/Flax 공부하기 전에
   1. JAX/Flax 소개와 예시
2. JAX 특징
   1. Numpy에서부터 JAX 시작하기
   2. JAX의 JIT 컴파일
   3. 자동 벡터화/자동 미분
   4. JAX의 난수
   5. pytree 사용하기
   6. JAX에서의 병렬처리
   7. 상태를 유지하는 연산
3. Flax 소개
   1. Flax CNN 튜토리얼
   2. 심화 튜토리얼
4. JAX/Flax 활용한 딥러닝 모델 만들기
   1. 순수 JAX로 구현한 CNN
   2. ResNet
   3. DCGAN
   4. CLIP
   5. DistilGPT2 미세조정 학습
5. TPU 환경 설정

결론

딥러닝 모델 학습에 막대한 리소스가 필요한 대규모 모델 연구에 깊이 관여 중이라면 JAX 를 고려해볼만 하다. 따라서 대규모 모델 연구 즉, LLM 같은 걸 연구하는 사람들에게 추천한다. 지난 약 1년 동안 나온 논문과 릴리스를 살펴보면 구글에서 많은 연구가 JAX로 옮겨갔다는 사실을 알 수 있다고 한다. 모든 것을 충족시키는 만능 프레임워크는 없다. 어떤 문제를 해결할 것인가, 배포할 모델의 크기, 플랫폼 등등에 따라 달라진다. 조그마한 모델을 학습시키기 위해 JAX를 고려할 필요는 없다는 의미이다. 그러나 대규모 연산의 확장성을 염두에 둔다면 JAX도 고려해볼만 하다. 

• 목차
  1. JAX/Flax를 공부하기 전에
     • JAX/Flax에 대한 소개와 예시
     • 함수형 프로그래밍에 대한 이해
     • JAX/Flax에서 자주 사용하는 파이썬 표준 라이브러리
     • JAX/Flax 설치 방법
  2. JAX의 특징
     • NumPy에서부터 JAX 시작하기
     • JAX의 JIT 컴파일
     • 자동 벡터화
     • 자동 미분
     • JAX의 난수
     • pytree 사용하기
     • JAX에서의 병렬처리
     • 상태를 유지하는 연산
  3. Flax 소개
     • Flax CNN 튜토리얼
     • 심화 튜토리얼
  4. JAX/Flax를 활용한 딥러닝 모델 만들기
     • 순수 JAX로 구현하는 CNN
     • ResNet
     • DCGAN
     • CLIP
     • DistilGPT2 미세조정 학습
  5. TPU 환경 설정
     • 코랩에서 TPU 설정하기
     • 캐글에서 TPU 세팅하기
     • TRC 프로그램 신청하기
• 총평
  • 모든 코드를 github에 공개하여 실습해볼 수 있음(notebook으로 공개되어서 바로 colab에서 실행해 볼 수 있음)
  • 매우 전문적인 서적으로 실제 프레임워크를 사용하여 ML/DL 프로젝트를 진행하는 사람들에게 빠르게 입문하는 것에 도움을 주는 책
  • ML/DL에 전문지식이 없는 사람에게는 내용의 벽이 너무 높음
  • JAX, Flax에 대한 소개(1~3장), 프레임워크를 이용한 실제 모델 훈련(4장)으로 구성되어 있어 기본적인 내용을 배우고 난 후, 간단한 프로젝트에 실제 적용을 해 보는 형태로 구성되어 있음
  • 4장의 튜토리얼의 경우, 각 노트북별로 사전에 설치하거나, 로드해야할 dataset 등이 해당 부분에서 설명하는 코드를 실행하기 전에 잘 나열되어 있어서 작은 chapter 별로 구분되어 있는 노트북을 실행하기만 하면 결과를 볼 수 있음
  • 책의 코드를 일일히 타이프하고 결과를 보기까지 수많은 에러에 직면하던 과거와는 달리, 너무나 에러없이 결과를 재현하면서 확인할 수 있어 만족스러웠음

• Chapter 1 : JAX, Flax의 정의와 기초
  • Jax 는 구글 딥마인드에서 사용하고 있는 고성능 딥러닝 프레임워크(파이토치, 텐서플로 같은 프레임워크임)
  • JAX의 장점은 GPU, TPU에서 Numpy프로그램을 컴파일하고 실행할 수 있게 해주며, XLA로 컴파일이 가능해서 파이토치의 그래프 방식보다 빠른 속도로 학습과 추론이 가능함
  • Flax는 JAX에 좀더 엔지니어가 편하게 사용할 수 있게 만든 프레임워크
  • JAX/Flax 설치 방법
    • 현재 리눅스와 macOS에만 설치 가능함
    • 구글 코랩에서도 사용 가능함
      • 단, 책의 예시는 에러가 발생(jax 버전 관련 오류로 보임)
      • 아래와 같이 고쳐야 성공(단, 런타임 유형을 TPU v2-8로 선택해야함)
        
      
      
• chapter 2 : JAX의 특징
  • 고성능 수치계산과 머신러닝에서 중요한 JAX의 특징 : jax.numpy, JIT compile, 의사 난수 생성 등
  • 장점
    • Numpy같은 중요 라이브러리와 유사하게 처리 가능하여 학습이 쉬움
    • Gradient 함수같은 것도 jax.grad(함수명) 이렇게 쉽게 처리 가능함
    • JIT compile을 통해 빠른 연산 수행 가능
    • 자동벡터화를 통해 병렬연산이 가능
    • numpy보다 우수한 유사난수 기능, python의 container를 유사하게 처리하는 pytree 함수 지원
    • 병렬처리를 지원하여 간단한 함수 몇개로 병렬처리를 지원할 수 있음
    • 상태를 유지하는 함수형 연산을 이용하면서도 변수의 상태를 update할 필요가 있을 때도 update하는 변수를 JIT Compile 외부에 두어 원하는 upate를 수행 가능함
  • 단점(주의점)
    • JAX변환은 함수형 프로그래밍의 원칙을 따르기에 비순수함수를 사용할 수는 있지만, 사용자가 원하는 동자을 하지 않을 수 있음을 명확히 알고 수행해야 함

• chapter 3 : Flax 소개
  • 텐서플로와 파이토치의 성공에는 각각의 프레임워크를 쉽고 간결하게 사용하기 위한 API의 역할이 컸으며, Flax도 JAX를 쉽고 간결하게 활용하도록 개발되었음
  • 특징
    • 안정성 : JAX에서 상태를 처리할 때보다 좀더 버그 발생이 완화됨
    • 제어기능 : 대부분의 신경망 프레임워크보다 좀더 세분화된 제어 기능과 표현력을 제공
    • 함수형API : JAX에서 사용되는 변환을 모듈에서 직접 재정의하여 사용할 수 있음
    • 코드 간결성 :
  • 튜토리얼 1
    • Flax의 Linen API를 사용하여 Convolution Neural Network를 구축하고 , MNIST 데이터셋에 대해서 이미지 분류 학습 진행
    • 필요패키지 설치 → 데이터 로드하기 → 모델정의와 초기화 → 메트릭 정의하기 → TrainState 초기화 →훈련스텝과 평가스텝 정의하기 → 모델 학습하기 → 모델 추론하기 순서로 함수를 작성하고 실행해 볼 수 있음
    • Colab에 notebook을 열고 차례차례 실행해 보면서 결과를 확인하였음
    • 실행결과 (직접 colab에서 돌려 봄). 글자 인식 정도의 문제는 쉽게 해결하고 있음
  • 튜토리얼 2
    • 배치 정규화 적용 : 데이터의 분포가 달라지며 학습성능이 떨어지는 것을 보완하고 정규화하여 문제해결하는 방법
    • 결과(실제 노트북을 돌려서 나온 결과, 도서의 결과와는 일부 다름)
  • Drop Out 적용, 학습률 스케줄링, 체크포인트 예제도 모두 colab에서 잘 실행되며 결과도 일치함

• Chapter 4 : JAX, Flax를 활용한 딥러닝 모델 만들기
  • JAX/Flax를 이용해서 CNN, ResNet, DCGAN, CLIP 모델을 구현하고 최적화하는 실습 진행
  • 숫자인식 : CNN(JAX로만 구현), ResNet
    • CNN
      • 패키지로드 → MNIST 데이터셋 로드 → 레이어구현 → 네트워크정의 → 훈련단계 → 평가단계 → 훈련 및 평가 → 모델추론하기 순서로 구현
      • 레이어 구현 : 완전연결레이어, 합성곱레이어, 활성화 함수, 폴링레이어, 평탄화, 직렬화 구현
    • ResNet
      • 필요한 패키리 로드 → MNIST 데이터셋 로드 → ResNet 모델 정의 → 메트릭 정의 → TrainState 초기화 → 학습방법 정의 → 학습 순서로 구현
    • 학습결과 : CNN, ResNet 모두 이미지 식별을 탁월하게 수행함
  • DCGAN
    • GAN(Generative Adversarial Network, 생성적 적대 신경망)의 일종인 DCGAN(Deep Convolution GAN)을 MNIST 데이터셋을 이용하여 이미지 생성
    • 필요한 패키지 로드 → MNIST 데이터셋 로드 → 생성자와 판별자 모델을 정의하고 초기화 → 학습방법 정의 → TrainState 초기화 → 학습 순서로 구현
    • 학습에 33분 54초가 소요됨(colab)
    • 최종 이미지 생성결과(실제 수행결과) : 알아보기 명확한 숫자도 있지만 여러이미지가 숫자인지 그림인지 확인하기 어려운 것도 상당수 임 → z_dim, epoch를 늘려보는 것을 권고함
  • CLIP(Contrastive Language-image Pre-training)
    • OpenAI에서 개발한 대규모 이미지-텍스트 쌍으로 학습된 멀티모달 신경망
    • 비지도 학습을 이용하며, 이미지-텍스트간 유사성을 최대화하고, 서로 다른 이미지-텍스트 쌍 간의 유사성을 최소화하가 위한 대조손실(contrastive loss) 사용
    • 학습결과 : 54분 걸려서 학습 완료