최적화를 위해 : Custom OP

성능을 떨어뜨리는 연산들은 직접 고쳐 보자.

Performance evaluation을 통해서 어떤 연산들이 성능을 떨어뜨리는지 발견했다면, 이 연산들을 제거하거나 변경하는 것도 성능 향상의 방법 중 하나가 될 수 있습니다.

Custom operation을 구현하는 것은 크게 두 가지 단계로 이루어져 있습니다.

1. Custom 연산자 구현 (C 혹은 Python)

2. 구현한 커스텀 연산자를 RKNN-Toolkit2를 사용하여 등록

   1. 이 과정에서 연산자의 이름, 입력/출력 텐서의 형태 등을 지정할 수 있다.

3. rknn 모델 export

   1. 커스텀 연산자가 포함된 RKNN 모델을 내보내기

1. Custom 연산자 구현하기

제가 다루는 Depth Anything 모델의 경우, exSDPAttention연산, resize연산에서 가장 시간이 많이 소요되었습니다.

12   exSDPAttention   FLOAT16  NPU    (6,64,1,1531),(6,64,1,1531),...          (6,64,1,1531)          198835/0/198835          46461        \                    100.0%/0.0%/0.0%     3444         exSDPAttention:/blocks.0/attn/MatMul_1_exsdpa
232  Resize           FLOAT16  CPU    (1,64,68,90),(1),(1),(4)                 (1,64,136,180)         0/0/0                    34261        \                    0.0%/0.0%/0.0%       765          Resize:/depth_head/refinenet2/Resize

따라서 우선 테스트로 exSDPAttention 연산을 Custom Operation으로 구현해 보도록 하겠습니다.

import numpy as np
from rknn.api import RKNN
from rknn.api.custom_op import get_node_attr

class exSDPAttention:
    op_type = 'exSDPAttention'

    def shape_infer(self, node, in_shapes, in_dtypes):
        out_shapes = in_shapes.copy()
        out_dtypes = in_dtypes.copy()
        return out_shapes, out_dtypes

    def compute(self, node, inputs):
        query, key, value = inputs[:3]
        
        attention_scores = np.matmul(query, key.transpose(-2, -1))
        attention_probs = np.softmax(attention_scores, axis=-1)
        context_layer = np.matmul(attention_probs, value)
        
        outputs = [context_layer]
        return outputs
    

if __name__ == '__main__':

    custom_model_path = 'depth_anything_vits_19.onnx'

    # Create RKNN object
    rknn = RKNN(verbose=True)

    # Pre-process config
    print('--> Config model')
    rknn.config(target_platform='rk3588')
    print('done')

    # Register cstSigmoid op
    print('--> Register exSDPAttention op')
    ret = rknn.reg_custom_op(exSDPAttention())
    if ret != 0:
        print('Register exSDPAttention op failed!')
        exit(ret)
    print('done')

    # Load model
    print('--> Loading model')
    ret = rknn.load_onnx(model=custom_model_path)
    if ret != 0:
        print('Load model failed!')
        exit(ret)
    print('done')

    # Build model
    print('--> Building model')
    ret = rknn.build(do_quantization=False)
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('done')

    # Export rknn model
    print('--> Export rknn model')
    ret = rknn.export_rknn('depth_anything_OP19_custom_op_ver_1.rknn')
    if ret != 0:
        print('Export rknn model failed!')
        exit(ret)
    print('done')

    rknn.release()

2. 성능 평가하기

그리고 이렇게 생성된 depth_anything_OP19_custom_op_ver_1.rknn을 Accuracy Analysis 기능을 이용해서 정확하지만 느렸던 기존의 Depth anything rknn 모델 (fp) 과 레이어별 출력을 비교합니다.

import os
import urllib
import traceback
import time
import sys
import numpy as np
import cv2
from rknn.api import RKNN
import urllib.request
from rknn.api.custom_op import get_node_attr


ONNX_MODEL = 'depth_anything_vits_19_jiwon_final.onnx'
DATASET_PATH = '/home/jiwon/Orange_Pi/RKNN/Projects/rknn_model_zoo/examples/pidnet/python/calibration_data/list.txt'

class exSDPAttention:
    op_type = 'exSDPAttention'

    def shape_infer(self, node, in_shapes, in_dtypes):
        out_shapes = in_shapes.copy()
        out_dtypes = in_dtypes.copy()
        return out_shapes, out_dtypes

    def compute(self, node, inputs):
        query, key, value = inputs[:3]
        
        attention_scores = np.matmul(query, key.transpose(-2, -1))
        attention_probs = np.softmax(attention_scores, axis=-1)
        context_layer = np.matmul(attention_probs, value)
        
        outputs = [context_layer]
        return outputs
    

if __name__ == '__main__':

    custom_model_path = 'depth_anything_vits_19_jiwon_final.onnx'

    # Create RKNN object
    rknn = RKNN(verbose=True)

    # Pre-process config
    print('--> Config model')
    rknn.config(
        mean_values=[0.485*255, 0.456*255, 0.406*255], 
        std_values=[0.229*255, 0.224*255, 0.225*255], 
        target_platform='rk3588',
    )
    print('done')

    # Register cstSigmoid op
    print('--> Register exSDPAttention op')
    ret = rknn.reg_custom_op(exSDPAttention())
    if ret != 0:
        print('Register exSDPAttention op failed!')
        exit(ret)
    print('done')

    # Load model
    print('--> Loading model')
    ret = rknn.load_onnx(model=ONNX_MODEL)
    if ret != 0:
        print('Load model failed!')
        exit(ret)
    print('done')

    # Build model
    print('--> Building model')
    ret = rknn.build(do_quantization=False)
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('done')

    # Accuracy analysis
    print('--> Accuracy analysis')
    ret = rknn.accuracy_analysis(inputs=['./2023-11-23_18-19-35_607986_front.jpg'], output_dir='./custom_op_snapshot',target='rk3588')
    if ret != 0:
        print('Accuracy analysis failed!')
        exit(ret)
    print('done')


    rknn.release()

그러면 이렇게 성능 평가 지표와 함께 결과를 볼 수 있습니다. 저의 경우에는 , Custom Op로 바꾸니까 성능이 형편없어졌네요.