딥러닝 모델을 학습할 때 메모리 부족은 숙명과도 같습니다. 특히 860GB에 달하는 방대한 이미지 데이터셋을 학습시키다 보면, 모델 구조가 좋건 나쁘건 메모리 한계와의 싸움이 시작됩니다.

오늘은 저의 실전 경험을 바탕으로, 메모리 부족 속에서도 끝내 학습을 완주할 수 있었던 최적화 기법과 트러블슈팅 노하우를 공유합니다. 이 글이 비슷한 문제를 겪는 분들께 실질적인 도움이 되길 바랍니다.

---

✅ 1단계. Lambda Layer 제거: 모델의 안정성부터 확보

• 문제: Lambda Layer 사용 시, TensorFlow/Keras 버전 차이로 TFLite 변환 실패, GPU 메모리 예측 불가
• 해결: Lambda 대신 **Custom Layer (예: PreprocessingLayer)**를 생성
• 장점:
  • 모델 직렬화 및 변환 과정에서의 안정성
  • 명확한 레이어 명명 및 구조 파악

class PreprocessingLayer(tf.keras.layers.Layer):
    def call(self, inputs):
        return applications.efficientnet.preprocess_input(inputs)

---

✅ 2단계. 모델 경량화: EfficientNetB3 → B0

• B3는 성능이 좋지만 파라미터 수가 많고 메모리 사용량이 높음
• EfficientNetB0로 교체하여 메모리 사용량을 획기적으로 절감
• 경량화된 네트워크에서도 Data Augmentation과 Regularization으로 충분한 성능 확보

---

✅ 3단계. 이미지 크기 & 배치 사이즈 최적화

설정	변경 전	변경 후
이미지 크기	300×300	224×224
배치 사이즈	32	8 → 16

• 이미지 크기 축소: 메모리 사용량을 제곱 수준으로 줄일 수 있음
• 배치 사이즈 조정: GPU OOM(Out-Of-Memory) 방지
• 실험적으로 최적의 조합을 찾아야 함 (절대적인 정답은 없음)

---

✅ 4단계. Mixed Precision Training (혼합 정밀도 학습)

• float32 대신 float16을 활용해 메모리 절약
• tf.keras.mixed_precision.set_global_policy('mixed_float16') 로 적용
• 모델 학습 속도와 GPU 메모리 사용량 모두 최적화

policy = tf.keras.mixed_precision.Policy('mixed_float16')
tf.keras.mixed_precision.set_global_policy(policy)

---

✅ 5단계. Dropout & Regularization 조절

• Dropout: 0.3 → 0.2로 점진적 감소
• 과도한 Dropout은 underfitting 유발, 적절한 조정이 필수
• 데이터 증강: RandomRotation, RandomZoom, RandomFlip 등의 증강 기법 적용으로 모델의 일반화 성능 확보

---

✅ 6단계. Optimizer 및 Callback 전략

• Optimizer: Adam + LossScaleOptimizer 로 혼합 정밀도에 맞게 최적화
• Callback:
  • EarlyStopping (patience=5~7)
  • ReduceLROnPlateau (patience=3)
  • ModelCheckpoint (best val_accuracy 저장)

이 조합으로 학습 안정성과 과적합 방지 효과를 극대화할 수 있었습니다.

---

✅ 7단계. Dataset 최적화 & 오류 무시

• .ignore_errors()로 문제가 있는 파일을 runtime에서 자동 스킵
• Dataset pipeline에서 float16 변환 & prefetch로 I/O 최적화

dataset = dataset.map(lambda x, y: (tf.cast(x, tf.float16), y))
dataset = dataset.prefetch(tf.data.AUTOTUNE)

---

✅ 8단계. 첫 배치 사전 테스트

훈련 전 첫 배치를 모델에 통과시켜 예측 shape 확인:

• 모델과 데이터셋의 shape mismatch 즉시 발견
• 불필요한 메모리 낭비 없이 오류를 사전에 방지

---

✅ 9단계. 모델 저장 후 용량 검증 & TFLite 변환

• 모델 크기가 너무 작으면 underfitting 가능성 확인
• TFLite 변환: float16으로 최적화
• 변환 후 용량 및 성능을 체크하여 실제 서빙 환경에서도 최적화된 성능 유지

converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16]

---

✅ 결론: 메모리 부족은 극복 가능하다

860GB라는 대용량 데이터도, 위의 단계를 거쳐가며

• 모델 안정화
• 메모리 절감
• 성능 최적화

를 달성할 수 있었습니다.

모델이 OOM으로 터질 때마다 절망하지 말고,
👉 이미지 크기, 배치 사이즈, precision, 모델 경량화
이 네 가지를 꼭 조정해보세요.

“메모리가 부족한 것이 아니라, 아직 최적화를 덜 한 것이다.”

이제 더 큰 데이터셋에도 당당히 도전할 수 있습니다. 🚀

---

📌 덤: 전체 최적화 프로세스 요약

1. Lambda 제거 → Custom Layer
2. EfficientNetB0 + Functional API
3. 이미지 & 배치 크기 조정
4. Mixed Precision 활성화
5. Dropout 등 Regularization 조정
6. Optimizer & Callback 세팅
7. Dataset: ignore_errors + prefetch
8. 첫 배치 테스트
9. 모델 크기 확인 + TFLite 변환

---