Thống nhất kỹ thuật triển khai mô hình Machine Learning với ONNX

Trong mỗi dự án ứng dụng AI, các kỹ sư AI thường sử dụng nhiều công cụ khác nhau tùy theo yêu cầu, mục đích hoặc thói quen để “chế tạo” AI model. Các AI model vì thế cũng đa dạng về công nghệ, định dạng và cả cách sử dụng, cách triển khai. Sự đa dạng này mang lại khả năng linh hoạt cho các kỹ sư AI, nhưng cũng là một thử thách cho quá trình ứng dụng AI vào hệ thống.

ONNX - Open Neural Network Exchange là một “định dạng chung” cho các mô hình Machine Learning. Có khả năng thống nhất kỹ thuật triển khai model, ONNX mang lại nhiều lợi ích trong việc tự động hóa, xây dựng quy trình xử lý cũng như tạo khả năng cập nhật liên tục cho hệ thống AI.

Để làm được điều này, ONNX và các công cụ hỗ trợ chúng ta 2 thao tác là chuyển đổi (convert) và sử dụng (consume/serving).

AI model thường được tạo ra bằng các framework, công cụ thông dụng như Keras, PyTorch, TensorFlow, SciKit-Learn… Mỗi framework, công cụ lại tạo ra định dạng mô hình AI khác nhau khiến cho việc sử dụng phải phụ thuộc vào công nghệ tạo ra chúng.

ONNX giải quyết vấn đề này bằng cách cung cấp nhiều phương tiện để chuyển đổi các định dạng khác nhau thành “.onnx”, từ đó tạo tiền đề cho khả năng thống nhất kỹ thuật serving. Việc chuyển đổi thành định dạng chung này giúp các kỹ sư AI có thể lựa chọn framework, công cụ để sử dụng mà vẫn có thể đẩy nhanh quy trình triển khai model.

ONNX model được sử dụng thông qua ONNX runtime. Đây là công cụ đa nền tảng giúp đẩy nhanh quá trình training và inference. ONNX runtime tương thích với nhiều deep learning framework thông dụng, giúp tận dụng phần cứng để thực thi AI model đơn giản hơn. Việc cài đặt và sử dụng ONNX runtime có thể được thực hiện nhanh chóng thông qua hướng dẫn này.

Để dễ hiểu hơn, chúng ta sẽ xem qua khả năng và tiện ích của ONNX khi được sử dụng làm “định dạng chung”. Đầu tiên, ta sẽ tìm hiểu về việc tạo ra và sử dụng model này thông qua SciKit-Learn và CNN. Sau đó, ta sẽ chuyển đổi các model này thành ONNX và sử dụng với ONNX runtime.

Đối với SciKit-Learn model:

Train script:

.......
pipeline = Pipeline([('tfidf', TfidfVectorizer()),('clf', naive_bayes.MultinomialNB())])
pipeline.fit(train_X,train_Y)

print(pipeline.predict_proba(['train_X_Tfidf']))

with open('./svm.pkl', 'wb') as handle:
    pickle.dump(pipeline, handle, protocol=pickle.HIGHEST_PROTOCOL)
.......

Đoạn script predict kết quả:

.......
with open('svm.pkl', 'rb') as file:
        sess = pickle.load(file)
data_input = json.loads(input)
preds = sess.predict_proba(data_input)
.......

Đối với CNN model:

Train script:

.......
cnn_model = model_class.create_model(vocab_size, tokenizer, w2v)
callbacks = model_class.get_callback()
cnn_model.fit(x_train, y_train,
                  batch_size=BATCH_SIZE,
                  epochs=EPOCHS,
                  validation_split=0.1,
                  verbose=1,
                  callbacks=callbacks)

model_class.visualize_result(cnn_model, x_test, y_test)
cnn_model.save('./cnn.h5')
.......

Đoạn script predict kết quả:

.......
#  sq_text is the result of preprocessing input data
sess = tf.keras.models.load_model('./cnn.h5')
input_name = sess.get_inputs()[0].name
pred_onnx = sess.run(None, {input_name: sq_text.astype(np.float32)})
.......

Như vậy, với mỗi model, ta cần bộ thư viện khác nhau (Keras/ SciKit-Learn) cho việc train và predict. Điều này đòi hỏi người sử dụng model phải có kiến thức của cả CNN và SciKit-Learn để có thể ứng dụng model một cách hiệu quả.

Để chuyển đổi thành ONNX, ta sử dụng các thư viện có sẵn như bên dưới.

Đoạn script chuyển đổi SciKit-Learn model thành ONNX:

.......
import skl2onnx
.......
onnx = skl2onnx.convert_sklearn(pipeline, name='naivebayesclf', initial_types=[('input', StringTensorType([1, 1]))])
skl2onnx.save_model(onnx, 'svm.onnx')
.......

Đoạn script chuyển đổi CNN model thành ONNX:

.......
import keras2onnx
.......
onnx_model = keras2onnx.convert_keras(cnn_model, cnn_model.name)
onnx.save_model(onnx_model, 'cnn.onnx')
.......

Sau đó, ta có thể dùng ONNX runtime dưới dạng thư viện để sử dụng trực tiếp model.

Đoạn script sử dụng ONNX model (chuyển đổi từ SciKit-Learn model):

.......
import onnxruntime as rt
.......
sess = rt.InferenceSession('svm.onnx')
label_name = sess.get_outputs()[1].name
pred_onnx = sess.run([label_name], input_feed={'input': np.array([data_input])})
.......

Đoạn script sử dụng ONNX model (chuyển đổi từ CNN model):

.......
import onnxruntime as rt
.......
sess = rt.InferenceSession('cnn.onnx')
pred_onnx = sess.run(None, {input_name: sq_text.astype(np.float32)})
.......

Ngoài ra, ta cũng có thể dùng ONNX runtime dưới dạng server thông qua docker image để host model thành API và sử dụng thông qua giao thức HTTP và gRPC.

Host ONNX model dưới dạng API server:

Đoạn script sử dụng HTTP endpoint:

def send(inputname, input):
    request_headers = {
        'Content-Type': 'application/json',
        'Accept': 'application/json'
    }

    PORT_NUMBER = 9001
    inference_url = 'http://127.0.0.1:' + \
        str(PORT_NUMBER) + '/v1/models/default/versions/1:predict'

    request_message = {
        'inputs': {
            inputname: {
                'dims': ['1', '300'],
                'dataType': 1,
                'floatData': input.tolist()
            }
        }
    }

    response = requests.post(
        inference_url, headers=request_headers, data=json.dumps(request_message))

    print(response.text)

Thông qua khả năng này của ONNX, ta có thể:

• Đơn giản hóa kỹ thuật sử dụng mà vẫn giữ được sự linh hoạt trong công nghệ tạo ra model
• Giảm độ phức tạp cũng như dung lượng của client sử dụng model
• Dễ dàng cập nhật model mà không cần thay đổi client code nhờ API server

Từ các đặc điểm trên, ta có thể ứng dụng ONNX để tự động hóa quá trình training, serving, và pre-train. Sử dụng ONNX giúp việc cập nhật model dễ dàng hơn. ONNX runtime có thể cài đặt trên nhiều loại thiết bị khác nhau, hỗ trợ nhiều ngôn ngữ khác nhau, giúp vệc triển khai ONNX model đa dạng, tối đa hóa khả năng tận dụng được nhiều loại phần cứng khi thực thi model.

Tuy vậy, việc sử dụng ONNX cũng có những thử thách riêng.

• Trong quá trình thử nghiệm, việc chuyển đổi thành ONNX model không phải lúc nào cũng ổn. Trong một thử nghiệm chuyển đổi SciKit-learn model có sử dụng tokenizer thành ONNX, mình đã gặp vấn đề về việc xuất ra model có dung lượng quá lớn (>2Gb), không phù hợp cho sử dụng.
• ONNX được hỗ trợ rất nhiều bộ chuyển đổi, nhưng mình vẫn chưa thử sử dụng để chuyển đổi tất cả các loại model. Việc sử dụng các bộ chuyển đổi này cũng yêu cầu phải có kiến thức nhất định về công nghệ AI đã dùng để tạo ra model.

ONNX giúp thu hẹp được khoảng cách giữa kỹ sư AI và các kỹ sư trong các lĩnh vực khác. Qua những trải nghiệm của bản thân, mình thấy ONNX hỗ trợ rất nhiều trong việc rút ngắn thời gian triển khai AI model vào các hệ thống và tự động hóa nhiều quy trình cần thiết.