OpenVino Training Kit for Scikit-Learn and Pytorch

README.md

@@ -15,6 +15,7 @@ This list gives an overview of all modules available inside the contrib reposito
 * [**Token Merging**](./modules/token_merging/): adaptation of [Token Merging method](https://arxiv.org/abs/2210.09461) for OpenVINO.
 * [**OpenVINO Code**](./modules/openvino_code): VSCode extension for AI code completion with OpenVINO.
 * [**Ollama-OpenVINO**](./modules/ollama_openvino): OpenVINO GenAI empowered Ollama which accelerate LLM on Intel platforms(including CPU, iGPU/dGPU, NPU).
+* [**ov_training_kit**](./modules/ov_training_kit): Training Kit Python library -- provides scikit-learn, PyTorch and Tensorflow wrappers for training, optimization, and deployment with OpenVINO on AI PCs.
 
 ## How to build OpenVINO with extra modules
 You can build OpenVINO, so it will include the modules from this repository. Contrib modules are under constant development and it is recommended to use them alongside the master branch or latest releases of OpenVINO.

modules/openvino_training_kit/.gitignore

@@ -0,0 +1,3 @@
+src/ov_training_kit.egg-info/
+dist/
+build/

modules/openvino_training_kit/README.md

@@ -0,0 +1,40 @@
+# OpenVino training kit
+
+Wrappers for scikit-learn and PyTorch models with OpenVINO optimization.
+
+## About
+
+This module provides easy-to-use wrappers for training, evaluating, and exporting classical (scikit-learn) and deep learning (PyTorch) models optimized for OpenVINO, targeting local AI PCs and edge deployment.
+
+
+## System Requirements
+
+- **Operating System:** Linux (Ubuntu 18.04+), Windows 10/11, Windows Server 2019+
+- **CPU:** x86-64 (Intel or AMD)
+- **Python:** 3.8, 3.9, 3.10, 3.11
+- **RAM:** 8GB+ recommended
+- **GPU:** Optional (not required)
+- **Note:** Intel Extension for PyTorch (IPEX) is only supported on Linux/Windows with x86-64 CPUs. On MacOS, some features may not be available.
+
+## Installation
+
+```bash
+pip install ov-training-kit
+```
+
+## Usage
+
+For detailed usage instructions and examples, please refer to the README files inside the `src/sklearn` and `src/pytorch` folders.
+
+---
+
+For questions, suggestions, or contributions, feel free to open an issue or pull
+
+## 🎓 Credits & License
+
+Developed as part of a GSoC
+
+### Authors
+
+- Leonardo Heim 
+- Shivam Basia 

modules/openvino_training_kit/setup.py

@@ -0,0 +1,53 @@
+from setuptools import setup, find_packages
+
+with open("README.md", "r", encoding="utf-8") as f:
+    long_description = f.read()
+
+setup(
+    name="ov_training_kit",
+    version="0.1.9",
+    description="Wrappers for scikit-learn and PyTorch models with OpenVINO optimization",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    url="https://github.com/openvinotoolkit/openvino_contrib",
+    packages=find_packages(where="src"),
+    package_dir={"": "src"},
+    include_package_data=True,
+    install_requires=[
+        "scikit-learn==1.2.2",
+        "scikit-learn-intelex==2023.1.1",
+        "torch>=1.12.0",
+        "openvino>=2023.0",
+        "nncf>=2.7.0",
+        "joblib>=1.2.0",
+        "numpy>=1.21.0,<2.0.0",
+        "psutil>=5.9.0",
+    ],
+    extras_require={
+        "ipex": [
+            "intel_extension_for_pytorch>=2.1.0"
+        ],
+        "dev": [
+            "pytest>=7.0.0",
+            "pytest-cov>=4.0.0",
+            "flake8>=6.0.0",
+            "black>=23.0.0",
+            "isort>=5.10.0",
+        ],
+        "docs": [
+            "sphinx>=5.0.0",
+            "sphinx_rtd_theme>=1.0.0",
+        ],
+    },
+    python_requires=">=3.8, <3.12",  
+    license="Apache-2.0",
+    classifiers=[
+        "Programming Language :: Python :: 3",
+        "License :: OSI Approved :: Apache Software License",
+        "Operating System :: OS Independent",
+        "Intended Audience :: Developers",
+        "Intended Audience :: Science/Research",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    ],
+     keywords="openvino scikit-learn pytorch machine-learning edge-ai",
+)

modules/openvino_training_kit/src/ov_training_kit/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (C) 2025 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+"""Root package for training_kit: exposes sklearn and pytorch wrappers."""
+
+from .sklearn import *
+from .pytorch import *
+
+__all__ = ["sklearn", "pytorch"]

modules/openvino_training_kit/src/ov_training_kit/pytorch/README.md

@@ -0,0 +1,91 @@
+# OpenVINO Kit - PyTorch Integration
+
+Wrappers for PyTorch models with OpenVINO for inference, quantization, and deployment.
+
+## Features
+
+- PyTorch model integration
+- Quantization-Aware Training (QAT) and mixed-precision (AMP) support
+- OpenVINO IR export and compilation
+- Built-in metrics for classification, regression, segmentation, and detection
+
+## Installation
+
+```bash
+pip install torch torchvision openvino nncf
+```
+
+## Basic Usage
+
+```python
+from torchvision.models import resnet18
+from ov_training_kit.pytorch import BaseWrapper
+
+model = resnet18(pretrained=True)
+wrapper = BaseWrapper(model)
+
+# Train
+from torch import nn, optim
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.Adam(model.parameters())
+wrapper.train(train_loader, criterion, optimizer, num_epochs=5, device="cuda")
+
+# Compile for OpenVINO IR (default)
+wrapper.compile()
+
+# Evaluate (default metric: accuracy for classification)
+def accuracy_metric(preds, targets):
+    return (preds.argmax(dim=1) == targets).float().mean().item()
+score = wrapper.evaluate(test_loader, accuracy_metric, device="cuda")
+print("Accuracy:", score)
+```
+
+## Metrics Examples
+
+**Classification**
+```python
+from ov_training_kit.pytorch import ClassificationWrapper
+classifier = ClassificationWrapper(model)
+acc = classifier.evaluate_accuracy(test_loader, device="cuda")
+```
+
+**Regression**
+```python
+from ov_training_kit.pytorch import RegressionWrapper
+regressor = RegressionWrapper(model)
+mse = regressor.evaluate_mse(test_loader, device="cuda")
+```
+
+**Segmentation**
+```python
+from ov_training_kit.pytorch import SegmentationWrapper
+segmenter = SegmentationWrapper(model)
+iou = segmenter.evaluate_iou(test_loader, num_classes=21, device="cuda")
+```
+
+**Detection**
+```python
+from ov_training_kit.pytorch import DetectionWrapper
+detector = DetectionWrapper(model)
+map_score = detector.evaluate_map(test_loader, metric_fn, device="cuda")
+```
+
+## Export to ONNX
+
+```python
+import torch
+from ov_training_kit.pytorch import export_model
+export_model(wrapper.model, input_sample=torch.randn(1, 3, 224, 224), export_path="model.onnx")
+```
+
+## Requirements
+
+- PyTorch >= 1.12
+- OpenVINO >= 2023.0
+- NNCF >= 2.7
+- Intel® Extension for PyTorch (IPEX) >= 2.1
+- Numpy
+
+## 🎓 Credits & License
+
+Developed as part of a GSoC

modules/openvino_training_kit/src/ov_training_kit/pytorch/__init__.py

@@ -0,0 +1,21 @@
+# Copyright (C) 2025 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+"""Pytorch models with OpenVINO optimizations"""
+
+from .base_wrapper import BaseWrapper
+from .classification_wrapper import ClassificationWrapper
+from .regression_wrapper import RegressionWrapper
+from .segmentation_wrapper import SegmentationWrapper
+from .detection_wrapper import DetectionWrapper
+from .compiler import compile_model
+
+
+__all__ = [
+    "BaseWrapper",
+    "ClassificationWrapper",
+    "RegressionWrapper",
+    "SegmentationWrapper",
+    "DetectionWrapper",
+    "compile_model",
+]