PyTorch Quantization tutorial : Static Quantization (CIFAR10)

개요

이전 글에 이어서 Quantization된 모델의 성능을 비교 하고자 한다.
이번에는 PyTorch의 Quantization기능을 중점으로 알아보고자 한다.

이 글을 쓰는 2022-03-07 기준 PyTorch의 stable release 버전인 1.10의 QUANTIZATION 문서를 참고했다.

Quantization 종류

Static Quantization

• Post Training Quantization 또는 PTQ 라고 함
• 학습된 모델을 양자화
• 가능하면 이전 레이어와 합침 (퓨즈, 예를 들면 Conv + Relu)
• 최적의 양자화 매개변수를 결정하려면 데이터 세트로 calibration해야 함
• PyTorch의 경우 현재 CPU에서만 지원함

Dynamic Quantization

• Weights는 미리 양자화, Activations은 추론 중에 동적으로 양자화
• 배치 크기가 작은 LSTM 및 Transformer 유형 모델에 적용
• 런타임에 관찰된 데이터 범위를 기반으로 활성화에 대한 스케일 팩터를 동적으로 결정
• 각 관찰된 데이터 세트에 대해 가능한 한 많은 신호가 보존되도록 스케일 팩터가 "tuned" 됨
• 실제 연산은 floating point = 속도의 이점이 크지 않음

아쉽게 Dynamic quantization은 Conv를 지원하지 않아 내가 관심 있는 Image 모델에 적용할 수 없다.
Dynamic quantization은 직접 확인하지 않았지만 속도 이점이 크지 않고 LSTM 같은데에 쓰인다고 한다.
그동안 내가 NPU 타겟으로 모델 경량화 한 방식은 Static Quantization인 것 같다.

Quantization Model 준비

• Quantization이 모듈 단위로 작동하기 때문에 Model을 Quantization 전에 모델 정의를 약간 수정해야 한다.
• 특히 모든 Quantization 기술에 대해 사용자는 다음을 수행해야 합니다.
  • 출력 재양자화가 필요한 작업을 기능에서 모듈 형식으로 변환
    • 예를 들어 torch.nn.functional.relu 대신 torch.nn.ReLU 사용
  • 하위 모듈에 .qconfig 속성을 할당하거나 qconfig_dict를 지정해서 양자화 부분 지정
    • 예를 들어 model.conv1.qconfig = None으로 설정하면 model.conv 레이어가 양자화되지 않음

Static Quantization을 위한 Model 준비

• 양자화/역양자화되는 위치 지정
  • Activation이 양자화/역양자화되는 위치를 QuantStub 및 DeQuantStub 으로 지정
• 양자화를 위한 특수 처리가 필요한 텐서 래핑
  • Torch.nn.quantized.FloatFunctional로 양자화를 위한 특수 처리가 필요한 텐서 래핑
  • 예를들어 출력 양자화 매개변수를 결정하기 위해 특별한 처리가 필요한 add 및 cat
• 퓨즈 모듈
  • 작업/모듈을 단일 모듈로 결합하여 더 높은 정확도와 성능을 얻을수 있음
  • 이것은 통합할 모듈 목록을 가져오는 torch.quantization.fuse_modules() API를 사용하여 수행됩니다. 현재 [Conv, Relu], [Conv, BatchNorm], [Conv, BatchNorm, Relu], [Linear, Relu] 융합을 지원함

Static Quantization

PyTorch + CIFAR10 모델 찾기

Static Quantization을 쓰기 위해서 Model의 forward() 함수를 수정해 양자화 될 지점을 수동으로 지정해야 한다. 기존에는 편의를 위해서 torch.hub에서 간단히 갖고 왔지만 모델 수정을 위해서 해당 모델의 torch.hub에 있는 github를 참조 해 코드를 수정해야 한다.

샘플에 사용될 모델은 ResNet18 모델이다. torchvision에서 제공하는 ImageNet으로 학습된 모델이 있지만, 캘리브레이션 시간이 오래 걸려서 CIFAR10 데이터셋으로 테스트 하고자 한다.

PyTorch + CIFAR10으로 미리 학습된 가중치 파일을 공유하는 오픈소스 Repo가 있어서 이를 참고 했다.

CIFAR10을 기준으로 여러 백본으로 학습된 파일을 구글 드라이브를 통해서 공유하고 있다. resnet18 기준으로 93.07%의 Val. Acc를 나타내고 있다고 한다.

해당 Repo에서 사용한 model 파일을 위 설명대로 적절히 수정해야 한다.

모델 수정

위 문서 설명에 따르면 add(+=) 사용을 위해서 nn.quantized.FloatFunctional 래핑을 해야 한다. 해당 부분을 찾아서 수정하면 된다. BasicBlock과 Bottleneck 두개 모두 수정하면 된다. 이 과정이 없으면 모델이 제대로 양자화 되지 않는다.

class BasicBlock(nn.Module):
  def __init__(...)
    ....
    self.skip_add = nn.quantized.FloatFunctional() # <- Add
    ...

  def forward(self, x: Tensor) -> Tensor:
    ....
    # out += identity
    out = self.skip_add.add(identity, out)
    out = self.relu2(out)

...

class Bottleneck(nn.Module):
  def __init__(...)
    self.skip_add = nn.quantized.FloatFunctional()
    self.relu2 = nn.ReLU(inplace=True)
    ...
  def forward(self, x: Tensor) -> Tensor:
    ...
    # out += identity
    out = self.skip_add.add(identity, out)
    out = self.relu2(out)

이제 모델을 로드하고 다운받은 가중치 파일을 로드 한다.

model_fp32 = resnet.resnet18(pretrained=False, num_classes=10)
state_dict = torch.load('resnet18.pt')
model_fp32.load_state_dict(state_dict)

양자화 전 원본 fp32 기준의 모델을 test dataset으로 Acc 확인하면 92.59가 나온다.

fp32_eval_loss, fp32_eval_accuracy = helper.evaluate_model(model=model_fp32, test_loader=test_loader, device=cuda_device, criterion=None)
print(fp32_eval_loss, fp32_eval_accuracy)

모델 Fuse

이제 이 모델을 fuse한다. 레이어간 최적화 할 수 있는 레이어는 합치는 것이다. 이 모델에서는 conv, bn, relu등을 합친다. 여기까지도 아직 fp32 기준이고 그저 모델을 합칠 뿐이다. 기존 모델을 복사해서 fuse 한다.

fused_model_fp32 = copy.deepcopy(model_fp32)
fused_model_fp32.eval()

fused_model_fp32 = torch.quantization.fuse_modules(fused_model_fp32, [["conv1", "bn1", "relu"]], inplace=True)
for module_name, module in fused_model_fp32.named_children():
    if "layer" in module_name:
        for basic_block_name, basic_block in module.named_children():
            torch.quantization.fuse_modules(basic_block, [["conv1", "bn1", "relu"], ["conv2", "bn2"]], inplace=True)
            for sub_block_name, sub_block in basic_block.named_children():
                if sub_block_name == "downsample":
                    torch.quantization.fuse_modules(sub_block, [["0", "1"]], inplace=True)

그리고 fuse가 잘 됐는지 확인 해 본다. 원본에 해당하는 model_fp32모델과 fuse한 모델인 fused_model_fp32의 결과가 같게 나오는지 확인 한다. 특별히 모델을 양자화 하거나 한게 아니기 때문데 1e-05 미만의 오차가 나와야 정상이다.

model_fp32.eval()
fused_model_fp32.eval()

assert helper.model_equivalence(model_1=model_fp32, model_2=fused_model_fp32, device=cpu_device, rtol=1e-05, atol=1e-05, num_tests=100, input_size=(1,3,32,32)), "Fused model is not equivalent to the original model!"

Quantization

이제 본격적으로 모델을 양자화 한다. 기존 모델은 fp32로 그냥 만들어진 모델인데, 이 기존 모델을 래핑하는 QuantizedResNet18이라는 모델을 만들고 기존 모델 전에 QuantStub(양자화), DeQuantStub(역양자화) 레이어를 추가했다. 양자화 되는 부분을 수동으로 정할 수 있는데 모델 전체로 정한 것이다.

class QuantizedResNet18(nn.Module):
    def __init__(self, model_fp32):
        super(QuantizedResNet18, self).__init__()

        self.quant = torch.quantization.QuantStub()
        self.model_fp32 = model_fp32
        self.dequant = torch.quantization.DeQuantStub()

    def forward(self, x):
        x = self.quant(x)
        x = self.model_fp32(x)
        x = self.dequant(x)
        return x

calibration을 위한 function을 하나 만든다. 그냥 모델에 데이터셋을 실행하는 것 뿐이다.

def calibrate_model(model, loader, device=torch.device("cpu:0")):

    model.to(device)
    model.eval()

    for inputs, labels in loader:
        inputs = inputs.to(device)
        labels = labels.to(device)
        _ = model(inputs)

실질적인 Quantization이 발생한다. 위에서 만든 QuantizedResNet18 모델을 만들고 fbgemm 모드로 설정하고 모델 양자화 준비 후에 Calibration을 한다. 시간이 오래 걸리기 때문에 GPU(cuda)로 하고, 양자화 자체는 또 CPU만 지원하기 때문에 cpu로 copy해서 Quantization 한다. convert하는데 시간이 좀 걸리긴 한다.
Calibration같은 경우 별거 안하는거 같지만 이걸 안하면 왠지 모르겠지만 Quantization했을 때 결과가 영 엉망으로 나온다. 개인적인 생각엔 성능이 조금 떨어질 수는 있어도 안해도 크게 상관 없을 것 같지만 안하면 아예 결과가 안나왔다.

quantized_model = QuantizedResNet18(model_fp32=fused_model_fp32)

quantized_model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
torch.quantization.prepare(quantized_model, inplace=True)

# Calibration
calibrate_model(model=quantized_model, loader=test_loader, device=cuda_device)
quantized_model = quantized_model.to(cpu_device)

quantized_model_int8 = torch.quantization.convert(quantized_model, inplace=True)
quantized_model_int8.eval()
print(quantized_model_int8)

이제 Quantization된 모델의 test dstaset의 acc를 확인하면 92.48정도가 나온다. 원본 모델의 acc가 92.59인걸 비교하면 0.1% 정도 떨어진 것을 확인 할 수 있다. 아주 Quantization이 잘 된 것 같다.

fp32_eval_loss, fp32_eval_accuracy = helper.evaluate_model(model=quantized_model_int8, test_loader=test_loader, device=cpu_device, criterion=None)
print(fp32_eval_loss, fp32_eval_accuracy)

Test

실재 Quantization된 모델에 이미지를 넣어서 돌려보면 그 결과가 아주 정확하게 나온다. 확실하게 고양이를 분류할 수 있다.

cat 0.9831167459487915
dog 0.004510259255766869
frog 0.0019906775560230017
deer 0.001859520678408444
bird 0.001859520678408444

모델 파일의 크기를 비교하면 원본 모델은 43M, 경량화된 모델은 11M로 약 1/4 정도 줄었다.

추론 시간을 비교하면 cpu 실행 기준으로 원본 fp32 모델은 0.0074sec, 양자화된 int8 모델은 0.0013sec로 약 5.7배 정도 성능 향상이 있었다.

다음 포스팅에서는 더 큰 데이터셋 ImageNet이나 다른 비전 모델에 적용 하거나, 더 좋은 성능을 내기 위해서 Aware Training에 대해서 알아보고자 한다.

참고 소스

• https://github.com/dankernel/pytorch-static-quantization

참고 자료

• https://pytorch.org/docs/stable/quantization.html#torch.quantization.quantize_dynamic
  https://pytorch.org/tutorials/recipes/recipes/dynamic_quantization.html
• https://raki-1203.github.io/boostcamp_ai_tech/week_18/day_82/02.-Packing-Ripped-Models-The-Quantization-Exercises/
• https://gaussian37.github.io/dl-pytorch-quantization/
• https://leimao.github.io/blog/PyTorch-Quantization-Aware-Training/
• https://pytorch.org/docs/stable/quantization.html

pytorch quantization tutorialpytorch quantization tutorialpytorch quantization tutorialpytorch quantization tutorialpytorch quantization tutorial