diff --git a/README.md b/README.md index 3a270ed..3d52aa9 100644 --- a/README.md +++ b/README.md @@ -165,6 +165,7 @@ results = trainer.fit() | [📖 API Reference](https://sonanceai.github.io/datamint-python-api/client_api.html) | ![Intermediate](https://img.shields.io/badge/level-intermediate-yellow) | Complete API documentation | | [🔥 PyTorch Integration](https://sonanceai.github.io/datamint-python-api/pytorch_integration.html) | ![Intermediate](https://img.shields.io/badge/level-intermediate-yellow) | ML workflow integration | | [🧠 Trainer Guide](https://sonanceai.github.io/datamint-python-api/trainer_api.html) | ![Intermediate](https://img.shields.io/badge/level-intermediate-yellow) | Built-in trainers, trainer lifecycle, and custom model integration | +| [🔍 Bringing an External Model into Datamint](https://sonanceai.github.io/datamint-python-api/inference.html) | ![Intermediate](https://img.shields.io/badge/level-intermediate-yellow) | Integrate, log, and deploy a model trained outside Datamint for inference through the UI | | [🛠️ Command Line Tools](https://sonanceai.github.io/datamint-python-api/command_line_tools.html) | ![Beginner](https://img.shields.io/badge/level-beginner-brightgreen) | Full reference for `datamint upload`, `datamint init`, and `datamint config` | | [🔒 SSL Troubleshooting](https://sonanceai.github.io/datamint-python-api/ssl_troubleshooting.html) | — | Fixing `SSLCertVerificationError` | | [📓 Notebooks](notebooks/) | ![Beginner](https://img.shields.io/badge/level-beginner-brightgreen) ![Intermediate](https://img.shields.io/badge/level-intermediate-yellow) ![Advanced](https://img.shields.io/badge/level-advanced-red) | Numbered, runnable tutorials. Start at `01_getting_started` and work through annotations, datasets, experiment tracking, deployment, and a full end-to-end example | diff --git a/docs/source/index.rst b/docs/source/index.rst index a65ccd8..e1dd035 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -67,6 +67,14 @@ For Developers Train models with built-in one-line trainers -- no training loop to write. + .. grid-item-card:: Bringing an External Model into Datamint + :link: inference + :link-type: doc + + :bdg-warning:`Intermediate` + + Integrate a model trained outside Datamint, log and register it in MLflow, and deploy it for inference through the UI. + .. grid-item-card:: Tutorials :link: tutorials :link-type: doc @@ -97,6 +105,7 @@ For Developers client_api pytorch_integration trainer_api + inference tutorials datamint_vs_raw_pytorch ssl_troubleshooting diff --git a/docs/source/inference.rst b/docs/source/inference.rst new file mode 100644 index 0000000..fc6d226 --- /dev/null +++ b/docs/source/inference.rst @@ -0,0 +1,314 @@ +.. _inference: + + +Bringing an External Model into Datamint +========================================== + +:bdg-warning:`Intermediate` + +This page is for models **trained entirely outside Datamint**, your own training loop, +Keras, Hugging Face, or any third-party framework. It walks through bringing such a +model in, logging it to MLflow, and deploying it for inference through the UI. + +Two ways to integrate, depending on what you have: + +- Your model fits a supported task family and you have a labeled Datamint project to + evaluate it against, and you want metrics logged for the team? Use the + :ref:`Shortcut ` below. +- Otherwise, any model, any task, no labeled data required? Use the + :ref:`Custom Adapter ` steps below. + +Both paths end at the same :ref:`Deploy ` step. + +If you want to **train** an external model using Datamint, see +:ref:`Training an External Model Through a Datamint Trainer `, +which covers swapping in your own architecture. + + +.. _external_model_shortcut: + +Shortcut: Your Model Fits a Supported Task Family +---------------------------------------------------- + +Wrap your pretrained weights as a Lightning module and call +``trainer.test(register_model=True)``. It runs zero training epochs (your weights are +untouched), computes test metrics, and registers the model in MLflow in one call. + +Useful when: + +- Your model can be expressed as a ``SegmentationModule``/``ClassificationModule`` + subclass (the same task-family shapes used by :doc:`trainer_api`). +- You already have a Datamint project with an annotated test split to evaluate against. +- You want metrics logged for the team alongside the model. + +.. code-block:: python + + import segmentation_models_pytorch as smp + import torch + + from datamint.lightning import SemanticSegmentation2DTrainer + from datamint.lightning.trainers.lightning_modules import SegmentationModule + + net = smp.UnetPlusPlus(encoder_name='resnet34', in_channels=3, classes=1) + net.load_state_dict(torch.load('my_checkpoint.pth', map_location='cpu')) + + + class ExternalSegModule(SegmentationModule): + def __init__(self, *args, **kwargs): + super().__init__(*args, class_names=['lesion'], **kwargs) + self.model = net + + def forward(self, x): + return self.model(x) + + + MODEL_NAME = "my-external-unet" + + trainer = SemanticSegmentation2DTrainer( + project="MyProject", + image_size=256, + model=ExternalSegModule, + model_name=MODEL_NAME, + ) + test_metrics = trainer.test(register_model=True) + +``test_metrics`` shows up in the Datamint dashboard +alongside metrics from any trainer-trained model. If you want to deploy your model, check +:ref:`Deploy `. Note that ``register_model=True`` does not set a ``champion`` alias, so use +``model_version=`` when deploying. + +.. _external_model_custom_adapter: + +Custom Adapter: Any Model, No Labeled Data Required +-------------------------------------------------------- + +The steps below build a generic adapter that works for any model, task, or dataset, +no metrics, no labeled project required. Use them when the shortcut above doesn't +apply. + +Load Your Checkpoint +~~~~~~~~~~~~~~~~~~~~~~ + +Load your model exactly as you would outside Datamint. + +.. code-block:: python + + import segmentation_models_pytorch as smp + import torch + + net = smp.UnetPlusPlus(encoder_name='resnet34', in_channels=3, classes=1) + state = torch.load('my_checkpoint.pth', map_location='cpu') + net.load_state_dict(state) + net.eval() + +If the checkpoint is a Lightning ``.ckpt`` produced by one of Datamint's own modules +(e.g. ``SegmentationModule``), use Lightning's own loader instead: +``SegmentationModule.load_from_checkpoint(path)``. + +Wrap It in a ``DatamintModel`` Adapter +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Datamint's prediction contract expects a list of +:py:class:`~datamint.entities.resource.BaseResource` objects in, and a +``list[list[Annotation]]`` out -- one annotation list per resource. Subclass +:py:class:`~datamint.mlflow.flavors.model.DatamintModel` and implement +``predict_default`` to bridge your model to that contract: + +.. code-block:: python + + import albumentations as A + import cv2 + import numpy as np + import torch + from albumentations.pytorch import ToTensorV2 + + from datamint.entities.annotations import ImageSegmentation + from datamint.mlflow.flavors.model import DatamintModel, ModelSettings + from datamint.mlflow.flavors.task_type import TaskType + + + class SegmentationAdapter(DatamintModel): + """Wraps a plain nn.Module for Datamint segmentation inference.""" + + task_type = TaskType.IMAGE_SEGMENTATION + + def __init__(self, torch_model, class_names, image_size=256, threshold=0.5, need_gpu=False): + super().__init__(torch_model=torch_model, settings=ModelSettings(need_gpu=need_gpu)) + self.class_names = class_names + self.image_size = image_size + self.threshold = threshold + self._transform = A.Compose([ + A.Resize(image_size, image_size), + A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), + ToTensorV2(), + ]) + + def predict_default(self, model_input, **kwargs): + device = self.inference_device + model = self.get_pytorch_model().to(device).eval() + + results = [] + for resource in model_input: + img = np.array(resource.fetch_file_data(auto_convert=True, use_cache=True)) + if img.ndim == 2: + img = np.stack([img, img, img], axis=-1) + elif img.shape[-1] == 4: + img = img[..., :3] + orig_h, orig_w = img.shape[:2] + + tensor = self._transform(image=img)['image'].unsqueeze(0).to(device) + with torch.inference_mode(): + logits = model(tensor) + probs = logits.sigmoid().squeeze(0).cpu().numpy() + + results.append([ + ImageSegmentation( + name=self.class_names[i], + segmentation_data=cv2.resize( + (probs[i] > self.threshold).astype(np.uint8), + (orig_w, orig_h), + interpolation=cv2.INTER_NEAREST, + ), + ) + for i in range(len(self.class_names)) + ]) + return results + + + adapter = SegmentationAdapter(torch_model=net, class_names=['lesion'], image_size=256) + +Smoke-Test the Adapter Locally +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Before logging anything to MLflow, call ``.predict()`` directly to make sure +``predict_default`` runs without errors: + +.. code-block:: python + + import io + from PIL import Image + from datamint.entities.resource import LocalResource + + buf = io.BytesIO() + Image.fromarray(np.random.randint(0, 255, (300, 400, 3), dtype=np.uint8)).save(buf, format='PNG') + dummy_resource = LocalResource(raw_data=buf.getvalue(), filename='dummy.png') + + predictions = adapter.predict([dummy_resource]) + for ann in predictions[0]: + print(f"{ann.name!r} mask shape={ann.mask.shape}") + + +Log & Register the Model in MLflow +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +``datamint.mlflow.flavors.log_model`` serialises your adapter (including the embedded +``nn.Module``) as an MLflow artifact, records ``task_type`` so the Datamint server knows +how to display predictions, and registers it in the Model Registry under a name you +choose. Calling ``datamint.mlflow.set_project`` first points MLflow at your Datamint +server and associates the run with the right project, so it shows up in the Datamint +dashboard. + +.. code-block:: python + + import mlflow + + import datamint.mlflow as datamint_mlflow + from datamint.mlflow.flavors import log_model + + PROJECT_NAME = "MyProject" + MODEL_NAME = "my-external-unet" + + datamint_mlflow.set_project(PROJECT_NAME) + mlflow.set_experiment(PROJECT_NAME) + + with mlflow.start_run(run_name='external_model_upload') as run: + mlflow.log_params({ + 'encoder': 'resnet34', + 'image_size': 256, + 'framework': 'segmentation_models_pytorch', + }) + model_info = log_model( + adapter, + task_type=TaskType.IMAGE_SEGMENTATION, + name='segmentation_model', + registered_model_name=MODEL_NAME, + ) + + print(f"Model URI : {model_info.model_uri}") + +Assign an Alias +^^^^^^^^^^^^^^^^^ + +Deployment resolves models by alias, not raw version number. Set one (commonly +``champion``) on the version you just registered: + +.. code-block:: python + + from mlflow import MlflowClient + + client = MlflowClient() + versions = client.search_model_versions(f"name='{MODEL_NAME}'") + latest_version = max(versions, key=lambda v: int(v.version)) + client.set_registered_model_alias(MODEL_NAME, 'champion', latest_version.version) + +At this point you can already verify the round trip by loading the model back and +predicting with it, exactly as you would for a trainer-registered model: + +.. code-block:: python + + from datamint import Api + + api = Api() + model = api.models.get_by_name(MODEL_NAME) + loaded_model = model.get_latest_version(alias='champion').load_model() + + resources = list(api.resources.get_list(project_name=PROJECT_NAME, limit=1)) + predictions = loaded_model.predict(resources) + + +.. _external_model_deploy: + +Deploy: Run Inference Through the UI +------------------------------------------ + +Deploying starts a serving instance for the registered model so predictions can be +triggered directly from the Datamint platform, without writing any code per prediction. +Resolve by ``model_alias`` if you set one (the custom adapter path above does this), or +by ``model_version`` directly if you came from the shortcut, which doesn't set an alias: + +.. code-block:: python + + job = api.deploy.start( + model_name=MODEL_NAME, + model_alias='champion', # or: model_version=1 + with_gpu=False, + ) + + import time + while True: + job = api.deploy.get_by_id(job.id) + if job.status in ('completed', 'failed', 'cancelled'): + break + time.sleep(15) + +Once ``job.status == 'completed'``, the model is available for inference from the +Datamint UI on any resource in the project. The same call is also available +programmatically, e.g. for batch/automated inference, via ``api.inference.submit``: + +.. code-block:: python + + inf_job = api.inference.submit( + model_name=MODEL_NAME, + model_alias='champion', + resource_id=resources[0].id, + ) + inf_job.wait() + predictions = inf_job.predictions + + +Related Examples +------------------ + +- `Deploying an Externally Trained Model `_ + -- the full runnable notebook this page is based on, including custom prediction + modes and updating a deployed model with a new version. diff --git a/docs/source/trainer_api.rst b/docs/source/trainer_api.rst index b6ae35f..5fdc47d 100644 --- a/docs/source/trainer_api.rst +++ b/docs/source/trainer_api.rst @@ -207,8 +207,10 @@ Any extra keyword arguments that are not consumed by Datamint are forwarded to ) -Using an External Model Inside a Datamint Trainer -------------------------------------------------- +.. _training_external_model: + +Training an External Model Through a Datamint Trainer +------------------------------------------------------- There are two supported patterns, and they are not equivalent. @@ -392,6 +394,12 @@ through YAML, independent of any live dataset object. The dataset itself is never serialized: ``load_from_file(path, dataset=...)`` requires a freshly-built dataset, mirroring the constructor's own requirement. +.. note:: + + Have a model that was trained **entirely outside Datamint** and want to integrate, + log, and deploy it for inference through the UI? See :doc:`inference` instead. + + Related Examples ---------------- diff --git a/notebooks/05_deployment/02_deploy_external_model.ipynb b/notebooks/05_deployment/02_deploy_external_model.ipynb index 22002e3..4b12e25 100644 --- a/notebooks/05_deployment/02_deploy_external_model.ipynb +++ b/notebooks/05_deployment/02_deploy_external_model.ipynb @@ -9,6 +9,8 @@ "\n", "This notebook shows how to take a model **trained outside of Datamint** — in your own script, a Jupyter notebook, or any third-party framework — and register, deploy, and run inference with it using Datamint's serving infrastructure.\n", "\n", + "> See [Bringing an External Model into Datamint](https://sonanceai.github.io/datamint-python-api/inference.html) for a narrative walkthrough of this same flow (integrate, log, deploy).\n", + "\n", "## When to use this tutorial\n", "\n", "Use this approach when you:\n",