Add ONNX optimizations component

.github/workflows/build_docker.yaml

@@ -16,7 +16,8 @@ jobs:
         directory: [
           data_prep,
           train,
-          eval
+          eval,
+          onnx_optimize
         ]
 
     steps:

.gitignore

@@ -4,6 +4,9 @@ torch_model
 onnx_model.onnx
 loss_plot.png
 local.py
+local.pth
+local.onnx
+local_2.onnx
 train_images
 val_images
 test_images
@@ -13,6 +16,11 @@ test_split.json
 f1_score_by_class.png
 precision_by_class.png
 recall_by_class.png
+train/comparison_image_*.png
+train/small_original.png
+train/small_transformed_custom.png
+train/small_transformed_weights.png
+optimized_model.onnx
 
 # pyenv
 .python-version

README.md

@@ -7,11 +7,6 @@ and [weather image classification dataset](https://www.kaggle.com/datasets/jehan
 study, the focus is on demonstrating key concepts like reproducibility, artifact tracking, and automation, rather than
 the specifics of tools and implementation details.
 
-Disclaimer — this project doesn't advocate for Kubeflow Pipelines as the definitive framework for developing machine
-learning workflows (alternative frameworks will be provided in the appendix), nor does it endorse Vertex AI as the
-optimal managed platform. The choice of these specific tools stems from pragmatic reasons — namely, their immediate
-availability in my existing setup or compatability with existing setup.
-
 The `ml-workflows` follows a `standard` project structure, as well as tooling, proposed in the
 repository [py-manage](https://github.com/martynas-subonis/py-manage).
 
@@ -50,7 +45,7 @@ training epochs, etc. The pipeline should guarantee reproducibility and ease of
 
 ![Pipeline Overview](readme_assets/overview.png)
 
-The [pipeline](template.py) consists of three main components:
+The [pipeline](template.py) consists of four main components:
 
 1. Data Preparation ([data_prep](data_prep)):
 
@@ -70,7 +65,8 @@ The [pipeline](template.py) consists of three main components:
 -
     - Fine-tunes the classifier head for the specific problem domain.
 -
-    - Outputs: trained PyTorch model, ONNX model, training metrics, and loss plot.
+    - Outputs: trained PyTorch model, ONNX model, ONNX model with transformations included as part of model graph, training metrics, and
+      loss plot.
 
 3. Model Evaluation ([eval](eval)):
 
@@ -81,6 +77,11 @@ The [pipeline](template.py) consists of three main components:
 -
     - Outputs: confusion matrix, weighted precision, recall, and F1-scores.
 
+4. ONNX Model Optimization ([onnx_optimize](onnx_optimize))
+
+-
+    - A simple component which loads ONNX model with transformations, optimizes it, and produces optimized model artifact.
+
 Each component has its own Docker image, ensuring a reproducible runtime environment through the use
 of [Poetry](https://python-poetry.org/docs/) and its `.lock` files. To guarantee consistency in results, all components
 employ fixed random seeds for randomized algorithms. Additionally, the training component includes
@@ -139,6 +140,7 @@ STAGING_BUCKET=  # Your GCS bucket, where Kubeflow Pipelines will persist its ar
 PREP_DATA_DOCKER_URI=  # Your URI of data_prep component Docker image. 
 TRAIN_MODEL_DOCKER_URI=  # Your URI of train component Docker image.
 EVAL_MODEL_DOCKER_URI=  # Your URI of eval component Docker image.
+ONNX_OPTIMIZE_DOCKER_URI = # Your URI of onnx_optimize component Docker image.
 ```
 
 - Have [Poetry installed](https://python-poetry.org/docs/#installation).

onnx_optimize/Dockerfile

@@ -0,0 +1,18 @@
+FROM python:3.12.5-slim AS builder
+
+RUN pip install --upgrade pip==24.2.0 && \
+    pip install poetry==1.8.3
+
+WORKDIR /app
+
+COPY pyproject.toml poetry.toml poetry.lock ./
+
+RUN poetry install
+
+FROM python:3.12.5-slim AS runtime
+
+WORKDIR /app
+
+ENV PATH="/app/.venv/bin:$PATH"
+
+COPY --from=builder /app/.venv .venv

onnx_optimize/component_func.py

@@ -0,0 +1,23 @@
+from kfp.dsl import Input, Metrics, Model, Output
+
+
+def onnx_optimize(
+    onnx_with_transform_model: Input[Model], optimization_metrics: Output[Metrics], optimized_onnx_with_transform_model: Output[Model]
+) -> None:
+    import logging
+    import time
+
+    import onnxruntime as rt
+
+    start_time = time.time()
+    logging.info("Started ONNX model optimization task.")
+
+    sess_options = rt.SessionOptions()
+    sess_options.graph_optimization_level = rt.GraphOptimizationLevel.ORT_ENABLE_ALL
+    sess_options.optimized_model_filepath = optimized_onnx_with_transform_model.path
+    rt.InferenceSession(f"{onnx_with_transform_model.path}.onnx", sess_options)
+    optimized_onnx_with_transform_model.framework = (
+        f"onnxruntime-{rt.__version__}, graphOptimizationLevel-{str(sess_options.graph_optimization_level)}"
+    )
+    optimization_metrics.log_metric("timeTakenSeconds", round(time.time() - start_time, 2))
+    logging.info("Successfully finished ONNX model optimization task.")