문제상황

pytorch 모델을 양자화 하기 위해서 static_quantize기법을 사용했었는데, 이는 이미 학습된 모델의 가중치만 불러와 양자화하는 방법이었다.

나는 모델 양자화를 고려하지 않고 학습시킨 모델을 불러와서 torch.quantization.convert를 수행했는데, 모델의 사이즈가 줄어들었고, 실제로 모델을 출력해봐도 quantized된 레이어로 잘 출력이 되었지만 output을 뽑으려고 할 때 마다 아래와 같은 오류를 마주쳤다.

import torch
import os
import argparse
import models.my_model as mm

def get_args():
    parser = argparse.ArgumentParser()
    # Data and model checkpoints directories
    parser.add_argument('--model_path', type=str, default="my_model_path", help='model.pth or model.pt')
    parser.add_argument('--model_dir', type=str, default="my_model_dir", help='../dir/dirr/dirrr/model_repository')
    
    args = parser.parse_args()
    return args

def print_model_size(model_path):
    print("%.2f MB" %(os.path.getsize(model_path)/1e6))

if __name__=="__main__":
    
    args = get_args()
    
    # Prepare input data
    input_data = torch.full((1, 3, 512, 512), 0.5)
    
    model_path = os.path.join(args.model_dir, args.model_path)
    save_path = os.path.join("./quantized_models", os.path.splitext(args.model_path)[0]+"_quantized.pt")
    model = mm.NestedUNet()
    model.load_state_dict(torch.load(model_path).state_dict())

	# Make predictions
    with torch.no_grad():
        output = model(input_data)
    print("-----before quntization-----")
    print_model_size(model_path) 
    print("max : ", torch.max(output))
    print("min : ", torch.min(output))
    print("mean : ", torch.mean(output))

    model_traced = torch.jit.trace(model, (input_data))

    backend = "qnnpack"
    model.qconfig = torch.quantization.get_default_qconfig(backend)
    torch.backends.quantized.engine = backend
    model_static_quantized_prepared = torch.quantization.prepare(model)
    model_static_quantized = torch.quantization.convert(model_static_quantized_prepared)

    model_static_quantized.eval()
    q_output = model_static_quantized(input_data)	# <- 이 부분에서 inference가 되지 않고 에러 발생

	print("-----after quantization-----")
    print_model_size(save_path) 
    print("max : ", torch.max(q_output))
    print("min : ", torch.min(q_output))
    print("mean : ", torch.mean(q_output))

NotImplementedError: Could not run 'quantized::conv2d.new' with arguments from the 'CPU' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build). If you are a Facebook employee using PyTorch on mobile, please visit https://fburl.com/ptmfixes for possible resolutions. 'quantized::conv2d.new' is only available for these backends: [QuantizedCPU, QuantizedCUDA, BackendSelect, Python, FuncTorchDynamicLayerBackMode, Functionalize, Named, Conjugate, Negative, ZeroTensor, ADInplaceOrView, AutogradOther, AutogradCPU, AutogradCUDA, AutogradXLA, AutogradMPS, AutogradXPU, AutogradHPU, AutogradLazy, AutogradMeta, Tracer, AutocastCPU, AutocastCUDA, FuncTorchBatched, FuncTorchVmapMode, Batched, VmapMode, FuncTorchGradWrapper, PythonTLSSnapshot, FuncTorchDynamicLayerFrontMode, PythonDispatcher].

원인추론

음... 일단 모델을 양자화 하면 모델 안에 있는 레이어들이 quantized 된 레이어로 바뀌는건 확인했는데... 에러를 해석해보면 대충 quantized된 레이어가 CPU나 CUDA에서 동작하지 않는다는 내용인거 같다. (위 상황에선 CPU인데 전부다 CUDA로 돌려서 해도 똑같이CUDA에서 못돌린다고 에러가 난다.)

양자화된 모델은 일반적인 eager mode에서 모델 돌리는거처럼 CPU나 CUDA에서 못돌리는건 알겠는데... 그럼 이걸 어떻게 돌려야하지? 싶어서 찾아보다가 양자화모델에서는 input도 양자화 해야하기 때문에 torch.quantization.QuantStub 모듈을 사용해서 입력을 양자화한 다음 넣어주어야 한다는 것을 발견했다.

q_input_data = torch.quantization.QuantStub()(input_data)
q_output = model_static_quantized(input_data)

그래서 위 와 같이 input_data를 양자화해서 넣어줬는데도 똑같은 에러가 났다. 여기서 더이상 뭘 해야 하는지 몰라 한참 해멨는데 공식문서랑 여기저기 자료를 한참 찾아보다가 알게된 점이 QuantStub()을 모두 모듈 안에 있는 init() 에서 해주고 있었다는 거였다.

나는 QuantStub()이 어떻게 동작하는지 모르고 그냥 단순히 입력을 양자화해주는 역할이겠거니 판단을 했다. 그래서 '모듈 내부에서 하나 외부에서 하나 별 차이없겠지' 싶어서 외부에서 QuantStub()을 적용 후 모델에 넣어줬는데, 이걸 모듈 안의 init에서 넣어줘야 제대로 동작하는것 같았다.

문제해결

class NestedUNetModel(nn.Module):
    def __init__(self, num_classes=29, input_channels=3):
        super().__init__()

        num_filter = [32, 64, 128, 256, 512]
        self.quant = torch.quantization.QuantStub()	# 입력을 양자화 하는 QuantStub()
        self.dequant = torch.quantization.DeQuantStub()	# 출력을 역양자화 하는 DeQuantStub()
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        # ....
        # ...
        # ..
        
    def forward(self, x):                    # (Batch, 3, 256, 256)
        x = self.quant(x)
        # ... model forward ...
        output = self.dequant(output)
        return output

그래서 모듈 내부에서 입출력을 양자화 및 역양자화 하도록 고친다음 다시 시도해보았더니 아래와 같은 결과가 나왔다.

NotImplementedError: Could not run 'aten::add.out' with arguments from the 'QuantizedCPU' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build). If you are a Facebook employee using PyTorch on mobile, please visit https://fburl.com/ptmfixes for possible resolutions. 'aten::add.out' is only available for these backends: [CPU, CUDA, Meta, MkldnnCPU, SparseCPU, SparseCUDA, SparseCsrCPU, SparseCsrCUDA, BackendSelect, Python, FuncTorchDynamicLayerBackMode, Functionalize, Named, Conjugate, Negative, ZeroTensor, ADInplaceOrView, AutogradOther, AutogradCPU, AutogradCUDA, AutogradHIP, AutogradXLA, AutogradMPS, AutogradIPU, AutogradXPU, AutogradHPU, AutogradVE, AutogradLazy, AutogradMeta, AutogradMTIA, AutogradPrivateUse1, AutogradPrivateUse2, AutogradPrivateUse3, AutogradNestedTensor, Tracer, AutocastCPU, AutocastCUDA, FuncTorchBatched, FuncTorchVmapMode, Batched, VmapMode, FuncTorchGradWrapper, PythonTLSSnapshot, FuncTorchDynamicLayerFrontMode, PythonDispatcher].

오... 같은 에런데 한가지 차이점이라면 quantized된 형식의 레이어가 CUDA나 CPU에서 사용하지 못한다는 에러였다면 이번에는 반대로 기존에 있던 add연산이 QuantizedCPU라는 백엔드에서 동작하지 못한다는 에러였다.

뭔가 quantized된 모델이 동작할 수 있는 환경을 만든것 같다!

근데 aten::add.out가 뭔지 몰라서 좀 찾아보니까 ArrayTensor의 약자이고 기본적인 Pytorch의 다차원 연산을 처리하는 역할을 수행하도록 pytorch의 백엔드 구현을 담당하는 모듈이라고 한다. 확실히 이전과는 반대로 양자화 모델이 동작가능한 백엔드를 구성한게 맞는거같다.

음... 추가로 뭔가를 더 해 줘야 하는거 같은데, 공식문서에서 아래와 같은 내용을 찾았다.

1. Specify where activations are quantized and de-quantized. This is done using QuantStub and DeQuantStub modules.

2. Use FloatFunctional to wrap tensor operations that require special handling for quantization into modules. Examples are operations like add and cat which require special handling to determine output quantization parameters.

3. Fuse modules: combine operations/modules into a single module to obtain higher accuracy and performance. This is done using the fuse_modules() API, which takes in lists of modules to be fused. We currently support the following fusions: [Conv, Relu], [Conv, BatchNorm], [Conv, BatchNorm, Relu], [Linear, Relu]

오 1번은 내가 방금 했던거고, 2번을 보니까 add나 cat 같은 연산을 FloatFunctional로 바꿔주어야 한다고 한다.

관련 문서 내용을 보니까

• add
• cat
• mul
• add_relu
• add_scalar
• mul_scalar

위 여섯가지 연산에 대해서 적용하면 된다고 하는데 원체 설명이 부실해서 정확히 무슨 동작을 하는 함수인지는 잘 모르겠다. 대략 파악한 바로는 양자화된 모델은 정수형 텐서를 사용하고, 정수 연산을 수행하기 때문에 양자화된 모델 안에서도 정확한 연산을 수행할 수 있도록 해주는 모듈 정도 인것 같다.

내 모듈에서는 torch.cat 연산을 사용하고 있었기 때문에 모두 FloatFunctional로 바꿔주었더니 잘 동작했다.

근데 QuantStub이랑 DeQuantStub이 정확히 뭘 하는거지?

위에서 QuantStub 이 어떻게 동작하는지 모르고 막 썼다가 에러가 나서 시간을 한참 썼는데, 굳이 모델 안에서 얘네를 정의해줘야 하는 이유가 궁금해서 좀 찾아보았다.

https://discuss.pytorch.org/t/what-do-de-quantstub-actually-do/82402

누군가 나랑 비슷한 생각을 질문한게 있어서 읽어봤는데 정리해보자면 모델을 양자화할 땐 기존 모델을 prepare 하고 convert하는 두 번의 과정을 거치는데

    model_static_quantized_prepared = torch.quantization.prepare(model)
    model_static_quantized = torch.quantization.convert(model_static_quantized_prepared)

1. prepare 단계에서는 출력텐서의 통계정보를 기록하기 위해 QuantStub의 출력에 대한 observer를 부여한다.

2. convert 단계에서는 QuantStub이 nnq.Quntize() 모듈로 교체되며 nnq.Quantize의 출력이 양자화 되어 입력으로 사용된다.

정도인거 같다. nnq.Quantize는 모델을 양자화하기위해 내부적으로 구현되어 있는 모듈같은데, 어쨌든 양자화 하기 위해 기존의 모델을 prepare하고 convert하는 과정에서 QuantStub이 사용되기 때문에 이를 모듈안에 넣어둬야 한다는 것을 알게 되었다!

QuantStub를 모듈안에 넣지않고 처음의 나처럼 밖에서 그냥 사용하게 되면 prepare하고 convert할 때 모듈 내에선 아무런 일도 일어나지 않기 때문에 모델의 양자화가 제대로 이루어지지 않는것이다.