Rohanberiwal commented on issue #41702:
URL: https://github.com/apache/airflow/issues/41702#issuecomment-2329345801
# ONNX Inference Operator for Apache Airflow
## Description
The `ONNXInferenceOperator` is a custom operator designed for running
inference using ONNX models within an Apache Airflow DAG. This operator
leverages the `onnxruntime` library to load an ONNX model and perform inference
on provided input data. The results of the inference are logged and returned.
### Components
1. **ONNXInferenceOperator**: A custom Airflow operator that initializes
with the path to the ONNX model and the input data. It performs inference in
the `execute` method and logs the results.
2. **run_onnx_inference**: A helper function that demonstrates how to run
inference using the `onnxruntime` library directly within a PythonOperator.
This function is provided as an alternative approach to using the custom
operator.
3. **DAG Definition**: Defines an Airflow DAG named `onnx_inference_dag`
that schedules the inference task to run once.
## Code
```python
import onnxruntime as ort
from airflow import DAG
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from airflow.operators.python import PythonOperator
from datetime import datetime
class ONNXInferenceOperator(BaseOperator):
@apply_defaults
def __init__(self, model_path: str, input_data: dict, *args, **kwargs):
super(ONNXInferenceOperator, self).__init__(*args, **kwargs)
self.model_path = model_path
self.input_data = input_data
def execute(self, context):
session = ort.InferenceSession(self.model_path)
input_name = session.get_inputs()[0].name
result = session.run(None, {input_name: self.input_data})
self.log.info(f"Inference result: {result}")
return result
def run_onnx_inference():
model_path = '/path/to/your/model.onnx'
session = ort.InferenceSession(model_path)
input_name = session.get_inputs()[0].name
input_data = {"your_input_key": [[1.0, 2.0, 3.0]]}
result = session.run(None, {input_name: input_data})
print(result)
with DAG(
dag_id='onnx_inference_dag',
start_date=datetime(2023, 1, 1),
schedule_interval='@once',
catchup=False
) as dag:
inference_task = ONNXInferenceOperator(
task_id='onnx_inference_task',
model_path='/path/to/your/model.onnx',
input_data={"your_input_key": [[1.0, 2.0, 3.0]]}
)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
