SBERT → ONNX 변환 및 Test example

들어가며..

지난 22년 7월 27일, ONNX 란 무엇인지에 관한 소개 글을 올렸었다. 해당 글은 본 블로그에서 항상 top-1 조회수를 차지했다. 필자 생각에는 많은 회사와 산업에서 인공지능 모델을 서빙하려 하기 때문에 많은 관심을 받았다고 생각한다. 이러한 관심에 힘입어 필자가 진행했던 Sentence-BERT 모델을 ONNX 변환시키는 간단한 예시를 포스팅하고자 한다.

---

 

Sentence-BERT 특징

ONNX 변환 전, Sentence-BERT의 Input/Output 형태를 알아야 한다. 형태는 다음과 같다.

• Input: text (ex. '나는 어바웃타임 영화를 좋아합니다.')
• Output: n차원 vector (ex. [0.1754, 0.7749, ...] 보통 768 차원 사용)

Input/Output 형태를 알아야 하는 이유를 간단하게 설명하면 ONNX 변환은 Process 전체 중 일부를 감싸주기 때문이다. 여기서 감싸준다는 말의 의미는 input이 들어갔을 때 감싸준 모든 연산을 마친 뒤 나오는 결과를 output으로 보여준다는 의미이다. 만약 해당 output 이후 추가 연산이 필요한 Process가 있다면 ONNX 연산과는 별개가 된다.
ONNX 변환 시 Model 연산을 포함한 전체 Process에서 어디서부터 어디까지의 연산을 감싸줄지 초기 설계를 잘 해야된다. 만약 해당 설계가 미비하면 ONNX 모델의 output을 다시 연산해야 하는 불필요한 상황이 발생하기 때문이다.

ONNX 변환 시 초기 설계의 중요성

Sentence-BERT는 문장을 vector로 변환하여 vector 들 간의 관계를 통해 정보를 추출하는 알고리즘이다. 때문에 Output으로 768차원의 vector가 나오게 된다. (일반적인 BERT의 hidden_size 차원 = 768 차원)
필자는 배포된 Ko-Sentence-BERT 모델(kosbert-sts)을 사용하여 추가적인 학습을 통해 프로젝트에 사용될 분류 task 과제를 수행했다. ONNX 변환까지 저작권 문제가 되지 않는 선에서 공유해보겠다.

---

 

예시 코드

필자가 개발한 소스 코드 중 저작권 문제가 될 수 있는 영역을 제외한 예시 코드를 공유한다.
문제가 될 수 있는 영역은 '~~~' 표시로 대신했다.

• torch 모델을 ONNX 변환하는 메서드: 'model2onnx'
• 변환된 ONNX 모델을 load하는 메서드: '_load_sbert_session'
• ONNX 모델 사용 Inference하는 메서드: '_encode'
• ONNX 모델 load ~ Inference하는 메서드: 'model_infer'

class ExportServModel(~~~):
    def __init__(
        self, 
        weight_path: Optional[str] = None,
        model_name: Optional[str] = None,
        output_path: Optional[str] = f'./results/export',
        device: Optional[str] = f'gpu',
        max_length: Optional[int] = 512,
        run_onnx: Optional[bool] = False
    ):
        """
        weight_path: load할 model_dir
        output_path: model_output 저장 경로 1
        model_name: model_output 저장 경로 2

        device: onnx model 추론 device 설정 (cpu or gpu)
        max_length: sentence max length
        run_onnx: onnx model을 사용한 inference

        model_output = output_path + model_name
        """
        super().__init__(
            output_path = output_path,
            model_name = model_name,
            weight_path = weight_path
        )

        if device == 'cpu':
            fast_onnxprovider = 'CPUExecutionProvider'
        elif device == 'gpu':
            fast_onnxprovider = 'CUDAExecutionProvider'
        else:
            assert fast_onnxprovider is None or fast_onnxprovider in ['cpu', 'gpu']

        self.max_length = max_length
        self.onnx_folder = output_path
        self.model_path = weight_path
        self.fast_onnxprovider = fast_onnxprovider
        self.export_model_name = os.path.join(self.model_output, f"model.onnx")
        self.sentences = [
            '오닉스 테스트 중 입니다.',
            '샘플 예시 입니다.',
            '아웃풋이 어떻게 나올까요?'
        ]

        if not os.path.exists(self.model_output):
            os.makedirs(self.model_output)

        if run_onnx == True:
            weight_path = os.path.dirname(glob(self.weight_path + '/**/pytorch_model.bin', recursive=True)[0])
            print(f'weight_path: {weight_path}')
            self.tokenizer = AutoTokenizer.from_pretrained(self.model_load_dir)

    def model2onnx(self):
        """
        BertModel → ONNX_Model 변환시키는 Convert 메서드
        """
        # hug/trans 라이브러리를 이용한 pretrained model and tokenizer 로드
        config_class, tokenizer_class, model_class = (AutoConfig, AutoTokenizer, AutoModel)

        weight_path = os.path.dirname(glob(self.weight_path + '/**/pytorch_model.bin', recursive=True)[0])
        print(f'weight_path: {weight_path}')

        config = config_class.from_pretrained(weight_path)
        tokenizer = tokenizer_class.from_pretrained(weight_path)
        model = model_class.from_pretrained(weight_path, config=config)        

        self.model = model

        model.eval()

        sample_args = {
            'padding': 'max_length',
            'max_length': self.max_length,
            'truncation': True,
            'return_tensors': 'pt'
        }

        sample_tensor = tokenizer('I love your eyes.', **sample_args)
        ### input shape 생성
        self.inputs = sample_tensor

        inputs = tuple(sample_tensor.values())
        input_names = list(sample_tensor.keys())

        with torch.no_grad():
            ### output shape 생성
            sample_output = model(**sample_tensor)

            self.outputs = {"last_hidden_state" : sample_output[0],
                            "pooler_output" : sample_output[1]}

            symbolic_names = {0: 'batch_size', 1: 'max_seq_len'}

            output_names=['last_hidden_state', 'pooler_output']
            dynamic_axes = {}
            for k, v in sample_tensor.items():
                dynamic_axes[k] = {0 : 'batch_size'}
            for name in output_names:
                dynamic_axes[name] = {0 : 'batch_size'}

            onnx_model = torch.onnx.export(
                model,
                args=inputs,
                f=self.export_model_name,
                export_params=True,
                input_names=input_names,
                output_names=output_names,
                dynamic_axes=dynamic_axes,
                opset_version=12
            )
        print(f"Model exported at: {self.export_model_name}")

    def _load_sbert_session(self):
        """
        ONNX 변환된 BertModel load 메서드
        """
        sess_options = onnxruntime.SessionOptions()
        sess_options.intra_op_num_threads = 1
        sess_options.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL

        session = onnxruntime.InferenceSession(
            self.export_model_name, sess_options, providers=[self.fast_onnxprovider])
        return session

    def _load_sbert_pooling(self):
        """
        Sentence-Transformers 라이브러리에서 Pooling module load 하는 메서드
        """
        model_json_path = os.path.join(self.weight_path, 'modules.json')
        with open(model_json_path) as fIn:
            modules_config = json.load(fIn)

        pooling_model_path = os.path.join(self.weight_path, modules_config[1].get('path'))
        pooling_model = Pooling.load(pooling_model_path)
        return pooling_model

    def _encode(self, session, pooling_model, sentences):
        """
        Pooling_model을 이용한 sentence embedding
        input: list 형식으로 된 sentence 묶음 또는 str 형식으로 된 sentence
        output: tensor 형식으로 된 sentence embedding 묶음
        """
        if isinstance(sentences, str):
            sentences = [sentences]

        inputs = self.tokenizer(
            sentences,
            padding=True,
            truncation=True,
            max_length=512,
            return_tensors="pt"
        )
        ort_inputs = {k: v.cpu().numpy() for k, v in inputs.items()}
        ort_outputs_gpu = session.run(None, ort_inputs)
        ort_result = pooling_model.forward(features={'token_embeddings': torch.Tensor(ort_outputs_gpu[0]),
                                                     'attention_mask': inputs.get('attention_mask')})
        result = ort_result.get('sentence_embedding')
        return result    

    def model_infer(self, sentences):
        """
        inference 할 때 사용하는 매서드.
        onnx model load한 뒤, Sentence-Transformers Pooling class 사용하여 sentence embedding 진행.
        """
        session = self._load_sbert_session()
        pooling_model = self._load_sbert_pooling()

        # inference using onnx model
        results = self._encode(
            session = session, 
            pooling_model = pooling_model, 
            sentences = sentences
        )
        return results

---

 

ONNX vs Torch 결과 비교

필자는 ONNX 변환 이후 vector 값을 활용해서 분류 task를 진행했었다.(디테일한 실험 내용은 저작권의 문제로 공개할 수 없다.)
ONNX 모델과 Torch 모델의 성능 및 속도를 비교한 결과를 대략적으로 공유한다.

1. 성능
   • ONNX 모델과 Torch 모델 간 성능 차이는 존재하지 않았다.
   • ONNX 모델을 사용하여 얻은 vector 값과 Torch 모델을 사용하여 얻은 vector 값을 비교한 결과,
     소수점 아래 5~7 자리에서 차이가 발생하였다. (그래서 성능 차이는 없었나보다.)
2. 추론 속도 (2가지 방법으로 비교)
   • 연산 방식 ⓐ: input text ~ 분류 결과 (total process)
     
     • Torch 모델이 ONNX 모델보다 추론 속도가 2.3배 빠르다. (on GPU)
     • ONNX 모델이 Torch 모델보다 추론 속도가 6.0배 빠르다. (on CPU)
   • 연산 방식 ⓑ: input text ~ vector (portion process)
     
     • ONNX 모델이 Torch 모델보다 추론 속도가 2.5배 빠르다. (on GPU)
     • ONNX 모델이 Torch 모델보다 추론 속도가 6.8배 빠르다. (on CPU)

추론 속도에서 이러한 결과를 보이는 이유는 '연산 방식 ⓐ  > on GPU' 에서 vector를 얻은 후의 process가 cpu로 계산되어졌기 때문이다. 물론 vector를 얻은 후 process를 gpu로 계산할 수 있지만 그럴바에야 ONNX 변환 시 감싸주는 영역을 total process로 하는 게 낫지 않을까?
만약 전체 Process를 ONNX로 변환하려면 모델의 구조를 변경할 필요가 있다. 모델 내 연산을 전체 Process로 변경한 뒤, ONNX 변환을 진행하면 GPU, CPU 모두 뛰어난 성능을 보일 것이다.

---

 

필자의 의견

실험을 통해 확인할 결과 ONNX 변환이 '잘' 된다면 (웬만하면) 모든 하드웨어에서 좋은 성능을 기대할 수 있었다. SBERT 뿐만 아니라 다른 NLP task 실험에서도 비슷한 결과를 보였다.

만약 CPU 환경과 GPU 환경의 결과가 다르다면 ONNX 변환을 어떻게 하였는지 체크해 보는 걸 추천한다.
다음 포스팅 계획은 필자 글 중 조회수 top-n에 해당하는 Continual Learning 실험 결과, Deep Learning Compiler 실험 결과를 공유하려 한다.