🧩 PyTorch로 배우는 양자화: QuantizableCNN 완전 정복!

AI 모델 경량화의 꽃, **양자화(Quantization)**를 코드로 직접 구현해보는 시간입니다! 👨‍💻👩‍💻 이론만 들으면 머리 아프니까, PyTorch 코드로 깔끔하게 정리해 드릴게요.

1️⃣ 핵심 개념: Stub과 Fusion 🤔

양자화를 하려면 모델에게 "여기서부터는 정수(INT8)로 계산해!"라고 알려줘야 해요. 그리고 계산 효율을 위해 여러 층을 하나로 합치는 퓨전(Fusion) 과정도 필요합니다.

• QuantStub: "자, 들어오는 실수(FP32) 데이터를 정수(INT8)로 바꿔!" (입구) 🚪
• DeQuantStub: "계산 끝났으니 다시 실수(FP32)로 바꿔서 내보내!" (출구) 🎁
• Fusion: Conv + BatchNorm + ReLU 3단계를 한 번에 계산해서 속도를 높이는 기술! ⚡️

2️⃣ QuantizableCNN 모델 코드 작성 📝

일반적인 CNN 모델에 딱 3가지만 추가하면 됩니다. 주석을 잘 봐주세요! 👀

Python

 

# @title QuantizableCNN.py
import torch
import torch.nn as nn
from torch.quantization import QuantStub, DeQuantStub

class QuantizableCNN(nn.Module):
    def __init__(self):
        super(QuantizableCNN, self).__init__()
        
        # 1. Stub 정의: 양자화의 시작과 끝을 담당해요 🏁
        self.quant = QuantStub()       # 입력을 FP32 -> INT8로 변환
        self.dequant = DeQuantStub()   # 출력을 INT8 -> FP32로 복구

        # 2. 일반적인 레이어 정의 (평범한 CNN과 똑같아요)
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(16)
        self.relu = nn.ReLU()
        
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(32)
        # relu는 재사용 가능

    def forward(self, x):
        # 3. 데이터 흐름 제어
        
        # [STEP A] 입구를 통과하며 INT8로 변신! 🤖
        x = self.quant(x)
        
        # [STEP B] 양자화된 상태로 연산 수행 (Conv -> BN -> ReLU)
        # 나중에 이 부분들은 하나로 합쳐질(Fuse) 거예요!
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu(x)

        # [STEP C] 출구를 통과하며 다시 FP32로 복귀! 🦋
        x = self.dequant(x)
        return x

    # 4. 퓨전(Fusion) 함수: 속도를 위해 레이어를 떡처럼 하나로 뭉칩니다! 🍡
    def fuse_model(self):
        # Conv + BN + ReLU를 한 번에 계산하도록 합체!
        torch.quantization.fuse_modules(self, [['conv1', 'bn1', 'relu'], ['conv2', 'bn2', 'relu']], inplace=True)

---

3️⃣ 양자화 실행 프로세스 (5단계 레시피) 🍳

코드를 짰다고 끝이 아니에요! 실제로 변환하는 과정이 필요합니다. 이를 **PTQ(Post Training Quantization)**라고 해요.

Python

 

# @title 양자화 실행 코드
# 1. 모델 인스턴스 생성 및 학습된 가중치 로드 (가정)
model = QuantizableCNN()
model.eval() # ⚠️ 중요: 반드시 평가 모드로 바꿔야 퓨전이 작동해요!

# 2. 레이어 퓨전 (Fusion) 🔥
# Conv, BN, ReLU가 따로 놀지 않고 하나로 합쳐집니다.
model.fuse_model()

# 3. 양자화 설정 (Configuration) ⚙️
# 'fbgemm'은 서버(x86)용, 'qnnpack'은 모바일(ARM)용 설정이에요.
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
torch.quantization.prepare(model, inplace=True)

# 4. 캘리브레이션 (Calibration) 📏
# 모델에게 실제 데이터를 조금 보여줘서 "아~ 데이터 범위가 -5에서 5 사이구나" 하고 
# Min/Max를 측정하게 합니다. (INT8로 변환할 기준을 잡는 과정)
# (실제로는 여기서 데이터로더를 돌려야 합니다)
# run_calibration(model, data_loader) 
print("데이터 범위 측정 완료! ✅")

# 5. 진짜 변환 (Convert) 🦋 -> 🐛
# FP32 모델이 INT8 모델로 완전히 바뀝니다!
torch.quantization.convert(model, inplace=True)

print("양자화 완료! 모델이 가벼워졌어요! 🎉")

---

💡 쉽고 정확한 비유: "해외여행 환전" ✈️💵

이 과정은 해외여행을 가서 돈을 쓰는 과정과 똑같아요!

1. QuantStub 👉 "인천공항 환전소" 🇰🇷➡️🇺🇸 : 한국 돈(FP32, 무겁고 정밀함)을 들고 가면 미국 여행을 못 하죠? 입국하자마자 달러(INT8, 가볍고 범용적)로 싹 환전합니다.
2. Conv, BN, ReLU (INT8 연산) 👉 "현지에서의 쇼핑" 🛍️ : 달러(INT8)를 가지고 물건을 사고 밥을 먹습니다. 이 과정은 현지 통화니까 아주 빠르고 계산이 쉽죠!
3. DeQuantStub 👉 "귀국 후 재환전" 🇺🇸➡️🇰🇷 : 여행이 끝나고 집에(User) 돌아올 때는, 남은 달러를 다시 한국 돈(FP32)으로 바꿔야 우리가 알아볼 수 있겠죠?
4. fuse_model 👉 "세트 메뉴 주문" 🍔🍟🥤 : 햄버거(Conv), 감튀(BN), 콜라(ReLU)를 따로 시키면 주문받는 데 오래 걸리죠? "여기 세트 하나요!"라고 묶어서 주문하면 주방에서 훨씬 빨리 나옵니다.