OpenMLCore - Official Documentation

Why choosing OpenMLCore ?

Models Ready to Use

Classic architectures customisable in few lines of code.

Modular Losses

Create custom loss functions based on DiceLoss, BCE and more, with weights on classes.

Pretrained weights

Direct access to pretrained weights to make the finetuning more efficient.

Utils

Handle datasets, metrics such as IoU, Accuracy and export your model in ONNX.

Quick Start

Choose your favorite method to install OpenMLCore.

Installation via pip

Install the latest stable version with PyPI :

pip install openmlcore

Cloner le Repository

To contribute or to fork the repo :

git clone https://github.com/openmlcore/openmlcore.git

Code Example

Multiclass Image Segmentation with U-Net, CustomLoss and export ONNX.

train_unet.py

import torch

from torch.utils.data import DataLoader, Subset
from torchvision.transforms import v2
from openmlcore.loss.loss import BCEWithLogitsLoss, CustomLoss, DiceLoss
from openmlcore.models.UNets import UNet
from openmlcore.base.data import SegmentationDataset
from openmlcore.metrics.metrics import *


# Normalization & Transforms:
transforms = v2.Compose([
    v2.ToImage(),
    v2.Resize((256, 256), interpolation=v2.InterpolationMode.NEAREST),
    v2.ToDtype(torch.float32, scale=True), 
    v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

NB_CLASSES = 5
# Data Preparation:
train_dataset = SegmentationDataset(root_dir="data/train", num_classes=NB_CLASSES, transform=transforms)
val_dataset   = SegmentationDataset(root_dir="data/val", num_classes=NB_CLASSES, transform=transforms)
train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True)
val_loader   = DataLoader(val_dataset, batch_size=4, shuffle=True)

# Metrics Configuration:
metrics = {
    "Accuracy": accuracy,
    "IoU": iou
}

weights_for_class = torch.tensor([1.0, 1.2, 1.5, 1.0, 1.0])

myUnet = UNet(
    input_channels=3, 
    output_channels=NB_CLASSES, 
    depth=4, 
    initial_filters=32, 
    loss_fn=DiceLoss(weights=weights_for_class), 
    metrics=metrics
)

myUnet.create_model()
myUnet.optimizer = torch.optim.Adam(params=myUnet.parameters(), lr=8e-4)

# Train the Model:
myUnet.train(train_loader, val_loader, epochs=3)
# Save Model & Export ONNX:
myUnet.save_model("unet_model.pt")
myUnet.save_model_in_onnx("unet_model.onnx", input_sample=torch.randn(1, 3, 64, 64))

Make the Machine Learning easy with OpenMLCore