[Vitis-AI] Quantizing the Model

이 글은 Xilinx의 Vitis AI User Guide(UG1414)를 따라 작성하였습니다.

Pruning과 Quantization Flow

Vitis AI Quantizer Workflow

먼저 pytorch를 사용하기 위해서 vai_q_pytorch를 설치해야 합니다. 작성일 기준 pytorch 1.2~1.10.2를 지원하지만 data parallelism은 지원하지 않습니다.

VAI quantizer는 floating-point model를 input으로 받아 전처리(batchnorm folding 및 inference에 필요하지 않은 node를 삭제)를 수행합니다.

 

Docker container를 이용한 vai_q_pytorch 설치

추가적인 설치 없이 이전에 설치한 docker 환경에서 CPU/GPU 컨테이너를 실행하고 Conda 환경을 활성화합니다.

conda activate vitis-ai-pytorch

 

다른 버전의 pytorch 사용을 위한 Conda 환경 복제 및 Source code를 이용한 설치 방법은 86, 87 page에 나와있습니다.  

 

Model Quantization 전 Float Model 검사(Inspect)

Float 모델을 양자화하기 전 어떤 operator가 어떤 장치(DPU/CPU)에서 실행될지 나타내는 파티션 정보를 출력합니다. DPU가 연산이 더 빠르기 때문에 최대한 DPU에서 실행하도록 하지만 할 수 없을 때 그 이유를 알려줍니다. 이를 이용해 최대한 DPU에서 모델이 작동할 수 있도록 하는데 도움을 얻을 수 있습니다.

Activation statistics를 얻고 양자화된 모델의 정확도를 높이기 위해 VAI가 여러번의 inference를 수행하며 activation들을 보정합니다. 여기에는 label이 필요하지 않고 약 100~1000개의 보정 이미지에서 잘 작동합니다. 보정 후 양자화된 모델은 DPU의 데이터 형식을 따르는 DPU deployable model(vai_q_pytorch의 경우 model.xmodel)로 변환됩니다. 이 모델은 Vitis AI compiler를 이용해 컴파일되고 DPU에 deploy 할 수 있습니다. Inspect를 수행하는 코드 및 순서는 아래 내용과 resnet18_quant.py 파일을 참고하시면 됩니다.

1. vai_q_pytorch 모듈을 import합니다.

from pytorch_nndct.apis import Inspector

2. 보드에서 사용하는 DPU의 이름, 또는 fingerprint로 inspector 생성합니다. ZCU104의 경우 "DPUCZDX8G_ISA1_B4096"를 사용합니다.

inspector = Inspector("0x603000b16013831") # by target fingerprint
or
inspector = Inspector("DPUCAHX8L_ISA0_SP") # by target name

3. float model을 검사합니다

input = torch.randn([batch_size, 3, 224, 224]) # model의 input shape에 맞는 형태로
inspector.inspect(model, input)

이제 이 파일을 아래와 같이 실행하시면 됩니다.

python resnet18_quant.py --quant_mode float --inspect

inspect가 성공적으로 끝나면 "./quantize_result" 폴더에 다음 3개의 파일이 생성됩니다.

1. inspect_{target}.txt: Target의 정보와 float model의 operation들에 대한 세부 내용
2. inspect_{target}.svg: image_format이 None이 아니면 생성된 inspection result를 시각화.
3. inspect_{target}.gv: image_format이 None이 아니면 생성된 inspection result의 dot source code 생성.

 

Model Quantization

vai_q_pytorch를 이용해 quantization을 수행하기 위해서는 다음 3가지가 필요합니다.

1. model.pth : pre-train 된 PyTorch 모델
2. model.py : float model(기존 모델)이 정의된 파이썬 스크립트 파일
3. calibration dataset : 약 100~1000장의 training dataset의 subset

Pytorch 모델이 quantizable 하게 만드려면 float model이 forward 방법만 포함해야하고 torch.jit.trace 함수를 이용한 jit trace test를 통과할 수 있어야합니다.

 

기존 float model script에 vai_q_pytorch API를 추가하여 quantization script 작성

pre-train 된 float model과 그 성능을 평가하는 코드가 있다면 quantizer API를 이용해 quantization 된 model로 성능 평가를 수행할 수 있습니다. Quantize calibration은 quant_mode flag를 "calib"로 설정하면 평가 단계에서 tensor의 quantization step을 결정합니다. 보정 후에는 quant_mode flag를 "test"로 설정해서 양자화된 model을 평가합니다.

1. vai_q_pytorch 모듈을 import 합니다.

from pytorch_nndct.apis import torch_quantizer, dump_xmodel

2. Quantization에 필요한 input과 quantizer를 생성하고 변환된 model을 얻습니다.

input = torch.randn([batch_size, 3, 224, 224]) # 이전과 동일하게 model의 input과 같게
quantizer = torch_quantizer(quant_mode, model, (input))
quant_model = quantizer.quant_model

3. 변환된 모듈을 forward 하여 평가 결과를 얻습니다.

acc1_gen, acc5_gen, loss_gen = evaluate(quant_model, val_loader, loss_fn)

여기서는 예제에서 정의한 대로 적은 것으로 기존 evaluation 코드에서 model만 변환된 model을 이용하고 input과 output은 각자의 evaluate 함수에 맞게 정하면 됩니다.

4. Quantization 결과를 얻고 model을 deploy 합니다.

if quant_mode == 'calib':
    quantizer.export_quant_config()
if deploy:
    quantizer.export_xmodel()
    quantizer.export_onnx_model()

 

 

Quantization 실행 및 결과 얻기

아래에서 계속 나올 quant_mode는 예제 코드를 보시면 각 mode에서 어떤 동작을 수행하는지 알 수 있습니다.

1. quant_mode는 calib로 하여 quantization을 수행합니다.

python resnet18_quant.py --quant_mode calib --subset_len 200

quantization 및 evaluation동안은 iteration 횟수를 제어하는 것이 중요합니다. 일반적으로 quantization에는 100~1000개의 data(예제에서는 image)면 충분하고, validation set은 전부 사용하면 됩니다. 여기서 iteration 횟수를 제어하기 위해 subset_len 값을 설정하는데 이는 선택사항으로 코드 내에서 개별적으로 설정하여도 됩니다.

위 작업이 끝나면 2개의 파일이 ./quantize_result 폴더에 생성됩니다. 

1. Model.py : 변환된 vai_q_pytorch 형식 모델입니다. 파일명은 설정에 따라 달라집니다.
2. Quant_info.json : 텐서의 quantization step이 저장된 파일로, quantize 된 모델을 평가하기 위해선 이 파일이 필요합니다.

2. quant_mode를 test로 하여 Quantize 완료된 model을 평가합니다.

python resnet18_quant.py --quant_mode test

3. batch size를 1로, 중복된 반복을 방지하기 위해 subset_len을 1으로 설정하고 deploy를 실행하여 컴파일에 필요한 xmodel(및 onnx 형식 모델)을 생성합니다.

python resnet18_quant.py --quant_mode test --subset_len 1 --batch_size=1 --deploy

다음 두 결과 파일이 ./quantize_result 폴더에 생성됩니다. 이를 이용해 FPGA에 deploy 할 수 있습니다.

1. Model_int.xmodel : deployed XIR format model
2. Model_int.onnx : deployed onnx format model

onnx는 (Open Neural Network Exchange)로 서로 다른 환경에서 생성된 모델들을 사용할 수 있도록하는 신경망 표준 모델 포멧입니다.

 

Inspect 결과에 없는 permute 작업이 추가되는 경우 

Inspect 단계에서 sink_stanspose와 같은 형태로 permute가 추가되어 DPU가 permute를 지원하지 않아 CPU로 할당된다는 내용이 뜰때가 있습니다. 이는 Pytorch('NCHW') 와 XIR('NHWC')간의 순서 호환 문제로 Quantizer가 순서를 맞춰주기위해 permute를 추가하는 경우입니다. 이는 XIR convolution 연산시 자동으로 변경시킨 이후 reshape이나 view 등의 작업으로 순서를 다시 원래대로 돌려야 할때 추가되는데, 이를 막기위해서 직접 permute를 추가해줄 수도 있습니다.

permute가 지원되지 않는 경우 transepose를 이용할수도 있습니다.

참고 : https://github.com/Xilinx/Vitis-AI/blob/v2.5/src/Vitis-AI-Quantizer/vai_q_pytorch/example/inspector_tutorial.ipynb