문제 발생
크로스 인코더(Cross-Encoder)는 바이 인코더(Bi-Encoder)와 다르게 질문(S1)과 답변(S2)을 함께 임베딩하고
그 유사도를 학습하는 방식이다.
일반적으로 속도 때문에 Bi-Encoder를 사용해서 미리 임베딩한 뒤 비교하는 방식을 사용했지만,
크로스 인코더를 사용하면 유사도 스코어가 많이 증가한다, (실제로 73 -> 83 까지 올라갔다.)
하지만 크로스 인코더의 단점은 바로 속도...
최대한 가벼운 모델을 사용해 속도를 최소화 했고 이미 많이 사용해서 쓰고있는 Onnx 모델로 변환하고 양자화를 사용해 모델 속도를 최소화 하려고했다. 하지만 Sentence-Transfomer 모듈에서 학습한 모델은 임베딩만 해주더라.
ONNX Runtime을 사용하였더니 output 값이 (1, 64, 256)의 차원으로 나왔다
내가 원하는건 0~1사이의 유사도 값인데 말이다.
원인
원인은 모델 저장 방식에 있었다.
Sentence-Transfomer 는 CrossEncoder로 학습시켜도 임베딩 모델을 저장하고 실제 Predict할때는
transformers 의 AutoModelForSequenceClassification를 사용했다.
AutoModelForSequenceClassification를 사용할 경우 뒤에
"Linear(in_features=256, out_features=1, bias=True)" 레이어가 마지막에 붙어
(1, 64, 256)의 값을 하나의 스칼라 값으로 바꿔준다.
해결
Onnx 모델에도 마지막에 Linear 레이어를 추가해서 결과값을 스칼라 값으로 바꿔주면 해결 되는 문제였다.
그럼 Sentence-Transfomer CrossEncoder의 아웃풋과 Onnx 모델의 아웃풋이 동일해질 것이다.
import transformers
import transformers.convert_graph_to_onnx as onnx_convert
from onnxruntime.quantization import quantize_dynamic, QuantType
# 파이프라인으로 분류모델로 변경
pipeline = transformers.pipeline(
"text-classification", model=model, tokenizer=tokenizer)
# CPU 전용 Onnx 모델로 변경
model = model.to('cpu')
onnx_convert.convert_pytorch(pipeline, opset=11, output=Path("cross_encoder.onnx"),
use_external_format=False)
# Onnx 모델 가중치 양자화
quantize_dynamic("cross_encoder.onnx", "cross_encoder_uint8.onnx",
weight_type=QuantType.QUInt8)
구글링과 챗GPT를 열심히 괴롭힌 결과 해결방법을 찾았다.
트렌스포머의 파이프라인을 정의한 이후 파이프라인을 Onnx 모델로 변환하는 것이다.
from transformers import AutoTokenizer
import onnxruntime as rt
sess = rt.InferenceSession(f"cross_uint8.onnx", providers=['CPUExecutionProvider'])
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
sentence = tokenizer([["안녕","오 안녕 반가워"]], padding=True, truncation='longest_first', return_tensors="pt", max_length=256)
# Tensor 형태를 numpy로 변경
input_feed = {
"input_ids": np.array(sentence['input_ids']),
"attention_mask": np.array(sentence['attention_mask']),
"token_type_ids": np.array(sentence['token_type_ids'])
}
print(sess.run(None, input_feed)[0]) # array([[4.698374]], dtype=float32)
["안녕","오 안녕 반가워"] 의 두 연결된 대화는 "4.698374" 라는 스칼라 값을 얻었다.
이제 이것에 Sigmoid 함수를 적용하면 0~1의 값을 얻을 수 있다.
def sigmoid(x):
return 1 / (1 +np.exp(-x))
sigmoid(sess.run(None, input_feed)[0]) # array([[0.99097216]], dtype=float32)
최종적으로 0.99097 의 값이 나왔다, 즉 "안녕" 다음에 "오 안녕 반가워"는 99% 자연스럽다고 볼 수 있다.
*[안녕 , 자동차]는 0.123의 값이 나왔다. 모두 변환 전 CrossEncoder 의 Predict 결과 값과 동일하다.
결론
크로스 인코더의 단점은 느린 속도다, 하지만 난 가벼운 모델과 양자화를 통해 속도를
7.1ms(0.0071s) 에서 803 us(0.000803s)로 8~9배 개선시켰고
그러면서도 높은 정확도를 유지했다.