베타리더 후기 viii
지은이 소개 x
JAX/Flax LAB 소개 xii
들어가며 xiii
이 책에 대하여 xiv

CHAPTER 1 JAX/Flax를 공부하기 전에 1

1.1 JAX/Flax에 대한 소개와 예시 1
__1.1.1 JAX란 1
__1.1.2 Flax란 2
__1.1.3 JAX로 이루어진 기타 프레임워크들 3
__1.1.4 JAX 프레임워크 사용 예시 3
1.2 함수형 프로그래밍에 대한 이해 5
__1.2.1 부수 효과와 순수 함수 5
__1.2.2 불변성과 순수 함수 7
__1.2.3 정리하며 8
1.3 JAX/Flax에서 자주 사용하는 파이썬 표준 라이브러리 9
__1.3.1 functools.partial() 10
__1.3.2 typing 모듈 12
__1.3.3 정리하며 13
1.4 JAX/Flax 설치 방법 14
__1.4.1 로컬에 JAX/Flax 설치하기 14
__1.4.2 코랩에서 TPU 사용하기 14

CHAPTER 2 JAX의 특징 17

2.1 NumPy에서부터 JAX 시작하기 18
__2.1.1 JAX와 NumPy 비교하기 18
__2.1.2 JAX에서 미분 계산하기 19
__2.1.3 손실 함수의 그레이디언트 계산하기 21
__2.1.4 손실 함수의 중간 과정 확인하기 22
__2.1.5 JAX의 함수형 언어적 특징 이해하기 23
__2.1.6 JAX로 간단한 학습 돌려보기 25
2.2 JAX의 JIT 컴파일 28
__2.2.1 JAX 변환 이해하기 29
__2.2.2 함수를 JIT 컴파일하기 32
__2.2.3 JIT 컴파일이 안 되는 경우 34
__2.2.4 JIT 컴파일과 캐싱 37
2.3 자동 벡터화 39
__2.3.1 수동으로 벡터화하기 39
__2.3.2 자동으로 벡터화하기 41
2.4 자동 미분 42
__2.4.1 고차 도함수 43
__2.4.2 그레이디언트 중지 46
__2.4.3 샘플당 그레이디언트 49
2.5 JAX의 난수 52
__2.5.1 NumPy의 난수 52
__2.5.2 JAX의 난수 56
2.6 pytree 사용하기 59
__2.6.1 pytree의 정의 60
__2.6.2 pytree 함수 사용법 61
2.7 JAX에서의 병렬처리 65
2.8 상태를 유지하는 연산 69
__2.8.1 상태에 대한 이해 69
__2.8.2 모델에 적용하기 72

CHAPTER 3 Flax 소개 77

3.1 Flax CNN 튜토리얼 79
__3.1.1 패키지 로드하기 79
__3.1.2 데이터 로드하기 80
__3.1.3 모델 정의와 초기화 81
__3.1.4 메트릭 정의하기 84
__3.1.5 TrainState 초기화 84
__3.1.6 훈련 스텝과 평가 스텝 정의하기 85
__3.1.7 모델 학습하기 87
__3.1.8 모델 추론하기 89
3.2 심화 튜토리얼 90
__3.2.1 배치 정규화 적용 91
__3.2.2 드롭아웃 적용 95
__3.2.3 학습률 스케줄링 98
__3.2.4 체크포인트 관리 103

CHAPTER 4 JAX/Flax를 활용한 딥러닝 모델 만들기 105

4.1 순수 JAX로 구현하는 CNN 106
__4.1.1 패키지 로드하기 107
__4.1.2 데이터 로드하기 108
__4.1.3 레이어 구현 108
__4.1.4 네트워크 정의하기 115
__4.1.5 학습 및 평가 준비 116
__4.1.6 학습 및 평가 118
__4.1.7 추론 121
4.2 ResNet 122
__4.2.1 패키지 로드하기 123
__4.2.2 데이터 로드하기 123
__4.2.3 모델 정의 및 초기화 124
__4.2.4 메트릭 정의하기 129
__4.2.5 TrainState 초기화 129
__4.2.6 훈련 스텝과 평가 스텝 정의하기 131
__4.2.7 모델 학습하기 132
__4.2.8 결과 시각화하기 135
4.3 DCGAN 136
__4.3.1 패키지 로드하기 136
__4.3.2 데이터 로드하기 137
__4.3.3 모델 정의 및 초기화 138
__4.3.4 학습 방법 정의하기 140
__4.3.5 TrainState 초기화 143
__4.3.6 모델 학습하기 145
__4.3.7 결과 시각화하기 146
4.4 CLIP 148
__4.4.1 CIFAR10 데이터셋으로 CLIP 미세조정 진행하기 150
__4.4.2 JAX로 만들어진 데이터셋 구축 클래스 150
__4.4.3 이미지 데이터 구축 함수 뜯어보기 151
__4.4.4 CLIP 모델 불러오기 154
__4.4.5 CLIP에 사용하기 위한 전처리 및 미세조정 155
__4.4.6 모델 학습에 필요한 함수 정의하기 156
__4.4.7 하이퍼파라미터 설정과 TrainState 구축하기 160
__4.4.8 모델 저장하고 체크포인트 만들기 161
__4.4.9 요약 클래스 만들기 162
__4.4.10 학습에 필요한 스텝 정의와 랜덤 인수 복제 163
__4.4.11 모델 학습하기와 모델 저장하기 163
4.5 DistilGPT2 미세조정 학습 166
__4.5.1 패키지 설치 167
__4.5.2 환경 설정 168
__4.5.3 토크나이저 학습 169
__4.5.4 데이터셋 전처리 171
__4.5.5 학습 및 평가 173
__4.5.6 추론 178

CHAPTER 5 TPU 환경 설정 181

5.1 코랩에서 TPU 설정하기 181
5.2 캐글에서 TPU 세팅하기 182
5.3 TRC 프로그램 신청하기 183

마무리하며 186
찾아보기 188