CNN 8-bit Quantization: Jacob et al. (2018) 구현해보기

Yeontachi·2025년 9월 2일
AIDeepLearningquantization

Vision을 위한 Quantization

목록 보기
3/10

CNN 8-bit Quantization: Jacob et al. (2018)를 직접 Python 코드로 구현을 해보자.

공통 유틸 함수 정리

양자화 수식은 아래와 같다.

r=S(q−Z)r = S(q-Z)
q=round(rS+Z)q = \text{round}(\frac{r}{S} + Z)

이를 구현하려면 가장 먼저 필요한 것은, 정수 범위, 스케일과 제로포인트를 계산하는 함수들이라 생각한다.

먼저 정수 범위를 정하는 함수는 아래와 같다.

def qrange(num_bits = 8, signed = True):
	if signed:  	# int8 대칭: -127..127
    	return -(2**(num_bits-1)-1), (2**(num_bits-1)-1)
    else: 			# uint8 비대칭: 0..255
    	return 0, (2**num_bits-1)

이 함수는 주어진 비트 폭에서 표현 가능한 정수 구간 [qmin,qmax][q_{\text{min}}, q_{\text{max}}]을 반환한다.
양자화된 정수 값 qq는 반드시 이 범위를 벗어나지 않아야 하므로, 클램프(clamp)에 필요한 기준이 된다.

보통 두 가지 모드로 사용한다.

  • signed = True →\to 대칭 int8: [−127,127][-127, 127]
    (일부 구현은 [−128,127][-128, 127]을 쓰지만, 대칭성과 overflow 방지를 위해 -127을 사용했다.)
  • signed = False →\to 비대칭 uint8: [0,255][0, 255]

직관적으로, 가중치는 대칭, 활성화는 비대칭을 쓰는 경우가 많다고 한다.

다음은 비대칭 스케일 계산 함수이다.

def get_scale_zero_point_asym(a, b, num_bits=8):
	qmin, qmax = qrange(num_bits=num_bits, signed=False)
    a, b = float(a), float(b)
    eps = 1e-12
    if b-a < eps:
    	return 1.0, 0
    S = (b-a) / (qmax - qmin)
    Z = round(qmin - a/S)
    Z = int(max(qmin, min(qmax, Z)))
    return S, Z

이 함수는 비대칭 양자화에서 스케일 SS과 제로포인트 ZZ를 계산한다.

SS와 ZZ의 공식은 다음과 같다.

S=b−aqmax−qmin,    Z=round(qmin−aS)S = \frac{b-a}{q_{\text{max}}-q_{\text{min}}}, \,\,\,\,Z = \text{round}(q_{\text{min}} - \frac{a}{S})

여기서 만약 텐서 값이 거의 상수라면, S=1.0,Z=0S = 1.0, Z = 0으로 안전 처리를 한다.

이렇게 하는 이유는, 활성화처럼 데이터가 [0,+max][0, +\text{max}] 범위로 치우친 경우, 단순 대칭으로 두면 정수 공간이 낭비가 된다. 따라서 비대칭 스케일을 적용해 0을 정확히 표현할 수 있도록 조정하는 것이다.

다음은 대칭 스케일 계산함수 이다.

def get_scale_zero_point_sym(max_abs, num_bits=8):
	qmin, qmax = qrange(num_bits=num=bits, signed =True) # -127..127
    S = (max_abs / qmax) if max_abs > 0 else 1.0
    Z = 0
    return S, Z

대칭 양자화에서는 [−max,max][-\text{max}, \text{max}] 범위에 맞춰 스케일을 계산한다.

공식은 아래와 같다.

S=max∣r∣127,   Z=0S = \frac{\text{max}|r|}{127}, \,\,\, Z = 0

Z=0Z =0으로 두기 때문에 계산이 단순해지고, 특히 가중치(Weights)에서 자주 사용된다.
채널마다 분포가 다르다면, per-channel 방식으로 각 채널마다 다른 SS를 계산하면 오차를 줄일 수 있다.

다음은 텐서 최소/최대 추정 함수이다.

def tensor_minmax(x, percentile=None, ch_axis=None):
	if ch_axis is None:
    	flat = x.reshape(-1)
        if percentile is None:
        	return flat.min().item(), flat.max().item()
        lo = torch.quantile(flat, percentile/100.0).item()
        hi = torch.quantile(flat, 1 - percentile/100.0).item()
        return lo, hi
    else:
    	reduce_dims = tuple(d for d in range(x.ndim) if d != ch_axis)
        if percentile is None:
        	lo = x.amin(dim=reduce_dims)
            hi = x.amax(dim=reduce_dims)
        else:
        	lo = x.movedim(ch_axis, 0).reshape(x.size(ch_axis), -1)
            hi = lo.clone()
            k_lo = max(0, int(percentile/100.0 * lo.size(1)) - 1)
            k_hi = min(lo.size(1)-1, int((1 - percentile/100.0) * lo.size(1)))
            lo, _ = lo.kthvalue(k_lo+1, dim=1)
            hi, _ = hi.kthvalue(k_hi+1, dim=1)
        return lo, hi

이 함수는 주어진 텐서의 최소값/최대값을 추정한다.

단순 min/max 대신, 퍼센타일(percentile)을 지정하면 아웃라이어를 무시할 수 있다.

예를 들어, percentile = 0.1 →\to 하위 0.1%와 상위 0.1% 값은 잘라내고 범위를 잡는다.

ch_axis=None이면 전체 텐서를 하나로 보고 per-tensor 통계를,
ch_axis가 지정되면 채널별로 별도의 min/max(per-channel 통계)를 계산한다.

이렇게 하면 아웃라이어 하나 때문에 스케일이 너무 커지는 문제를 막을 수 있어, 양자화 정밀도를 높이는 효과가 있다.

Fake Quantization + STE

CNN 8-bit Quantization: Jacob et al. (2018) 이 글에서 설명했 듯이, 논문에서 PTQ 방식의 문제점을 해결하고자, QAT 방법을 제안했다.

QAT의 핵심 아이디어는 모델이 학습할 때부터 양자화 환경을 흉내 내는 것이다.

Forward, 일반적인 학습은 모든 연산이 FP32(32비트 부동소수점)로 진행된다. 하지만 QAT에서는 순전파(Forward pass) 단계에 가짜 양자화 연산(fake quantization)을 삽입한다.

가짜 양자화는 아래 과정을 거치게 된다.

1. Quantization(양자화)

q=round(xS)+Zq = \text{round}(\frac{x}{S}) + Z

실수 입력 x를 스테일 S로 나누고, 제로포인트 Z를 더해 정수로 변환한다.

q = torch.round(x / S) + Z

2. Clamp(포화 연산)
q가 표현 가능한 정수 범위 [qmin,qmax][q_{\text{min}}, q_{\text{max}}]를 벗어나면 가장 가까운 경계값으로 잘라낸다.

q_clamped = torch.clamp(q, qmin, qmax)

정수 표현 범위를 벗어난 값은 포화(saturation) 시킨다.
예를 들어, int8 대칭 : [−127,127][-127, 127], uint8 비대칭 : [0,255][0, 255]

x_min = (qmin - Z) * S
x_max = (qmax - Z) * S
mask = (x >= x_min) & (x <= x_max)

실수 도메인에서의 유효 구간을 계산한다. 정수로 클램프되는 경계값을 실수축으로 환산한 것이 x_min, x_max이며, mask는 클램프 없이 표현 가능한 입력 구간을 나타낸다. mask는 나중에 STE 게이팅에 사용된다.

ctx.save_for_backward(mask)
ctx.S = S; ctx.Z = Z; ctx.qmin = qmin; ctx.qmax = qmax

ctx는 forward와 backward 사이에 데이터를 전달하는 임시 저장소 역할을 한다.
ctx.save_for_backward(mask) 이 함수는 텐서를 안전하게 저장해 backward에서 그대로 불러올 수 있게 한다.(내부적으로 Pytorch는 ctx.save_for_backward로 저장된 텐서만 메모리에 효율적으로 캐싱해둔다.)

여기서 mask만 저장했는데 이유는 backward에서 gradient를 어디서 통과시킬지 판단할 때 필요하기 때문이다.(아래 예시: backward에서의 사용예시)

(mask,) = ctx.saved_tensors
grad_input = grad_output * mask

ctx.S = S; ctx.Z = Z; ctx.qmin = qmin; ctx.qmax = qmax는 디버그용으로 속성으로 값을 저장한 것이다.

3. Dequantization

x^=S⋅(q−Z)\hat{x} = S\cdot(q-Z)

정수를 다시 실수로 복원한다.

x_hat = S * (q_clamped - Z)
return x_hat

x_hat은 FP32지만, 값은 정수 격자에 맞춰진 상태(=가짜 양자화된 텐서)이다.

결과적으로 출력은 여전히 FP32 텐서이지만, 값은 정수 격자에 맞춰진 상태로 반환된다. 이렇게 하면 모델은 실제 INT8 환경에서 연산한 것처럼 양자화 오차(quantization noise)와 포화 효과(saturation)를 경험할 수 있다.

Backward, 문제는 라운딩(round)과 클램프(Clamp)는 미분이 불가능하다는 점이다. 역전파(Backward pass)에서 미분이 0이 되어버리면, 학습이 아예 진행되지 못한다.

이를 해결하기 위해 사용하는 기법이 바로 Straight-Through Estimator(STE, 직선 통과 추정자)이다. Bengio et al.(2013) 논문에서 제안된 방식으로, 아이디어는 단순하다.

  • Forward에서는 실제로 라운딩을 수행한다.
  • Backward에서는 라운딩의 미분을 항등함수처럼 근사한다.

즉, 수식으로 표현하면,

forward:y=round(x),   backward:∂y∂x≈1forward: y = \text{round}(x), \,\,\,backward:\frac{∂y}{∂x} ≈ 1

여기에 clamp까지 고려하면,
유효 범위 안에 있는 값들은 기울기를 그대로 통과시키고,
범위를 벗어나 포화된 값들은 기울기를 0으로 잘라낸다.

(mask,) = ctx.saved_tensors
grad_input = grad_output * mask.to(grad_output.dtype)
return grad_input, None, None, None, None

Fake Quantization + STE 전체 코드

class QuantDequantSTE(Function):
    @staticmethod
    def forward(ctx, x, S, Z, qmin, qmax):
        """
        x: FP32 텐서
        S: scale (broadcastable)
        Z: zero-point (broadcastable, 정수 취급)
        qmin/qmax: 정수 범위 (스칼라 int)
        """
        # 양자화
        q = torch.round(x / S) + Z
        q_clamped = torch.clamp(q, qmin, qmax)
        # 포워드 마스크: clamp 영역 안에 있는지(=기울기 통과)
        x_min = (qmin - Z) * S
        x_max = (qmax - Z) * S
        mask = (x >= x_min) & (x <= x_max)
        ctx.save_for_backward(mask)
        ctx.S = S
        ctx.Z = Z
        ctx.qmin = qmin
        ctx.qmax = qmax
        # 디양자화
        x_hat = S * (q_clamped - Z)
        return x_hat

    @staticmethod
    def backward(ctx, grad_output):
        """
        STE: clamp 범위 내부는 grad 그대로 통과, 밖은 0으로 컷.
        (S, Z, qmin, qmax는 학습하지 않으므로 None)
        """
        (mask,) = ctx.saved_tensors
        grad_input = grad_output * mask.to(grad_output.dtype)
        return grad_input, None, None, None, None

Test

QuantDequantSTE Module Test

아래 테스트 코드는, 전체 네트워크에 들어가기 전에 QuanDequantSTE 모듈 자체가 제대로 동작하는지 확인하는 테스트이다.

if __name__ == "__main__":
    torch.manual_seed(0)

    # 임의의 FP32 텐서 생성
    x = torch.randn(8, 16, 32, 32, requires_grad=True)

    # 스케일/제로포인트 계산 (비대칭, per-tensor)
    a, b = tensor_minmax(x)
    S, Z = get_scale_zero_point_asym(a, b)

    # 가짜 양자화 실행
    y = QuantDequantSTE.apply(x, S, Z, 0, 255)
    print("[FakeQuant] input range:", (float(x.min()), float(x.max())))
    print("[FakeQuant] output range:", (float(y.min()), float(y.max())))

    # grad 확인
    loss = y.sum()
    loss.backward()
    print("[FakeQuant] grad mean:", x.grad.abs().mean().item())

출력은 아래와 같다.

첫번째 출력 input range는 원래 텐서 x 값 범위를 나타낸다. 즉, 배치 내에서 최소 -4.34 ~ 최대 4.56까지 실수 값이 들어 있었다는 의미이다.

두 번쨰 출력 output range는 Fake Quantization 이후 y값의 범위이다. 결과를 비교해보면, 입력 범위랑 거의 같은데, 소수점 끝자리 정도만 달라진 것을 확인할 수 있다. 이유는, round(x/S) 때문에 값들이 정수 격자에 스냅되었기 때문이다. 결과적으로 아주 작은 오차만 생겼음을 보여준다.

마지막 출력 grad mean은 loss=y.sum() 후 역전파로 x.grad의 평균 절댓값을 찍은 값이다.

값이 거의 1.0에 가까운 건 STE가 잘 적용돼서 gradient가 끊기지 않고 그대로 통과했음을 의미한다. 만약 STE 없이 단순 round/clamp만 있었다면, grad가 다 0이 되어버려 학습이 불가능했을 것이다.

Simple QAT Test

아래 테스트 코드는 실제로 학습이 잘 되는지, QAT가 어떻게 작동하는지 보여주는 미니 실험이다.

class QATLinear(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.fc = nn.Linear(in_features, out_features)
        # weight용 scale/zero-point 추정을 위해 min/max 러닝 값 추적 가능

    def forward(self, x):
        # weight quant
        w = self.fc.weight
        w_min, w_max = w.min(), w.max()
        S_w, Z_w = get_scale_zero_point_sym(max_abs=max(abs(float(w_min)), float(w_max)))
        w_q = QuantDequantSTE.apply(w, S_w, Z_w, -127, 127)

        # input quant
        a, b = tensor_minmax(x)
        S_x, Z_x = get_scale_zero_point_asym(a, b)
        x_q = QuantDequantSTE.apply(x, S_x, Z_x, 0, 255)

        # 선형 연산
        y = F.linear(x_q, w_q, self.fc.bias)

        # output quant (per-tensor 비대칭)
        a, b = tensor_minmax(y)
        S_y, Z_y = get_scale_zero_point_asym(a, b)
        y_q = QuantDequantSTE.apply(y, S_y, Z_y, 0, 255)
        return y_q

# 실습용 미니 학습
if __name__ == "__main__":
    torch.manual_seed(0)
    model = QATLinear(10, 2)
    opt = torch.optim.SGD(model.parameters(), lr=0.1)

    for step in range(5):
        x = torch.randn(4, 10)
        target = torch.randint(0, 2, (4,))
        out = model(x)
        loss = F.cross_entropy(out, target)
        opt.zero_grad()
        loss.backward()
        opt.step()
        print(f"[Step {step}] loss={loss.item():.4f}")

위 코드에서 QAT가 적용된 선형 레이어를 만드는 방식은 다음과 같다.

class QATLinear(nn.Module):
    def forward(self, x):
        # 입력 x 양자화
        # 가중치 W 양자화
        # 출력 y 양자화
        return y_q

Forward 안에서 입력/가중치/출력을 모두 QuanDequantSTE로 감싼다. 즉, 학습 중 매번 가짜 양자화(Fake Quantization)를 흉내내는 구조이다.(훈련 중에도 정수 격자에 스냅된 값으로 연산을 시뮬레이션 한다.)

만약 QuantDequantSTE가 잘못 구현됐다면, loss가 변화없이 그대로 멈추거나 NaN이여야 하지만 출력이 잘 나오는 것으로 보아, gradient가 정상적으로 통과되고 학습이 이루어지고 있음을 확인할 수 있다.

Resnet18

실제 네트워크에서도 FakeQuant를 적용했을 때도 forward가 잘 동작하는지 확인해보자.

# ResNet18에 테스트해보기
from torchvision import models

if __name__ == "__main__":
    model = models.resnet18(weights=None, num_classes=10)

    # conv 레이어에 FakeQuant 래퍼를 적용해볼 수 있음
    # 예: conv1 레이어 테스트
    conv1 = model.conv1
    a, b = tensor_minmax(conv1.weight)
    S, Z = get_scale_zero_point_sym(max_abs=max(abs(float(a)), float(b)))
    w_q = QuantDequantSTE.apply(conv1.weight, S, Z, -127, 127)
    print("[ResNet18 conv1] 원래 weight 범위:", (conv1.weight.min().item(), conv1.weight.max().item()))
    print("[ResNet18 conv1] 양자화 weight 범위:", (w_q.min().item(), w_q.max().item()))

출력 결과는 다음과 같다.

원래의 weight 범위를 보면, 대략, (-0.0900, 0.0967)정도의 값을 갖고있다. 이런 값들은 학습된 네트워크의 필터 가중치 분포를 반영한다.

QuanDequantSTE를 통해 INT8 격자에 스냅후 다시 FP32로 복원한 범위는 다음과 같다. (-0.0898, 0.0967) 원래 범위와 거의 동일하지만, 소수점 몇 자리 차이가 보인다. 이 작은 차이가 바로 양자화 오차(quantization error)이다.

Forward 시점에서 가중치가 정수 격자에 맞춰졌다가 복원되었음을 보여준다. 값의 분포가 크게 달라지지 않았으므로, INT8 격자화가 안정적으로 적용된 것으로 보인다. 만약 잘못 구현됐다면 범위가 심하게 틀어지거나, 값이 전부 특정 값으로 뭉쳤을 것이다.

profile
Yeontachi
기초를 다지는 중입니다.📚🧑‍💻
이전 포스트

CNN 8-bit Quantization: Jacob et al. (2018)

다음 포스트

HAWQ: Hessian AWare Quantization of Neural Networks with Mixed-Precision (2019) Paper Review

0개의 댓글