Pytorch Quantization 적용

딥러닝 모델이 실제 device에 deploy 하는데 2가지 문제점이 있다.

느린 inference time
큰 model parameter size

Pytorch는 float32에서 int8로 데이터 크기를 줄여 연산을 하는 Quantization을 제공한다. 직접 짠 모델에서 quantization을 어떻게 적용하는지 알아보자. 전체코드는 이 링크에 있다.

Work Flow

float32에서 학습시킨 model 혹은 pretrain model을 가져온다.
model을 eveluation으로 변경 후 layer fusion을 적용한다. (Conv + BN + RELU)
forward의 input엔 torch.quantization.QuantStub(), output엔 torch.quantization.DeQuantStub()을 적용한다.
quantization configuration을 적용한다.
layer를 fuse한다.
QAT (quantized aware training) 을 진행한다.
모델을 cpu에 올려놓고 eval mode로 바꾼 후 float32모델을 int8모델로 변환시킨다.

Code

1. Declare Model

모델은 resnet 사용하기로 한다. 그리고 편의를 위해 학습은 미리 시켰다고 가정한다. 먼저 resnet의 BottleNeck을 선언하고 resnet18을 구현한다.

2,3. Deploy Layer Fusion

이후 각 모듈에 layer fusion을 적용한다. layer를 건드리지 않고 상속을 쓰면 결과적으로 parameter가 같기 때문에 QuantizableResNet18은 ResNet18의 파라미터를 쓸 수 있다. QuantizableBottleNeck에서는 두 tensor를 더하는 연산이 있으므로 기존의 방식이 아닌 FloatFunctional을 이용해야한다.

6,7. QAT, Convert int8 model

QAT를 진행한 다음에 float32모델을 int8모델로 변환시킨다.

Result

간단하게 MNIST dataset으로 학습시켰다. epoch 5, image size 224, optimzer Adam을 사용하였다.

	Accuracy (%)	Inference Time (ms)	Model Parameter (MB)
ResNet	98.79	66	46.31
Quantizable ResNet	96.42	33	11.69