Onnx Runtime Training (ORT)

pip install -r odt_requirements.txt

from onnxruntime.training import artifacts
from daiedge_vlab import dAIEdgeVLabAPI, OnDeviceTrainingConfig
import numpy as np
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import torch.nn as nn
import onnx
import torch
import torch.nn.functional as F

############################################
# Configuration
############################################

SETUP_FILE = "setup.yaml"

TARGET = 'rpi5'
RUNTIME = 'ort'
MODEL = 'artifacts.zip'

TRAIN_NB_SAMPLES = 10000
TEST_NB_SAMPLES = 2000

RESULT_DIR = "./results"

# ############################################
# # Step 1: Load, Preprocess MNIST, and Save Data
# ############################################

def prepare_datasets():
    # Load MNIST dataset.
    transform = transforms.ToTensor()
    train_dataset = datasets.MNIST(root=".", train=True, download=True, transform=transform)
    test_dataset  = datasets.MNIST(root=".", train=False, download=True, transform=transform)

    # Convert datasets to numpy arrays
    x_train = np.stack([np.array(img).squeeze() for img, _ in train_dataset])
    y_train = np.array([label for _, label in train_dataset])
    x_test  = np.stack([np.array(img).squeeze() for img, _ in test_dataset])
    y_test  = np.array([label for _, label in test_dataset])

    # Preprocess: Convert images to float32, normalize to [0,1], and add channel dimension.
    x_train = (x_train.astype(np.float32))[..., np.newaxis]
    x_test  = (x_test.astype(np.float32))[..., np.newaxis]

    # Reduce the size of the dataset for faster training (optional).
    x_train = x_train[:TRAIN_NB_SAMPLES]  
    y_train = y_train[:TRAIN_NB_SAMPLES]
    x_test  = x_test[:TEST_NB_SAMPLES]    
    y_test  = y_test[:TEST_NB_SAMPLES]

    # Preprocess labels.
    y_train = y_train.astype(np.int64)
    y_test = y_test.astype(np.int64)

    # Save preprocessed data as binary files.
    x_train.tofile("mnist_train_input.bin")
    y_train.tofile("mnist_train_target.bin")
    x_test.tofile("mnist_test_input.bin")
    y_test.tofile("mnist_test_target.bin")

    print(len(x_train), len(y_train), len(x_test), len(y_test))
    print("Preprocessed train/test data saved as .bin files.")
    return (x_train, y_train), (x_test, y_test)

train_set, test_set = prepare_datasets()

############################################
# Step 2: Define a Simple Keras Model.
############################################

class CNNModel(nn.Module):
    def __init__(self):
        super(CNNModel, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3)  # input channels = 1, output = 32
        self.pool1 = nn.MaxPool2d(kernel_size=2)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3)
        self.pool2 = nn.MaxPool2d(kernel_size=2)
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(64 * 5 * 5, 128)  # 28x28 → 26x26 → 13x13 → 11x11 → 5x5
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))  # [B, 32, 26, 26]
        x = self.pool1(x)          # [B, 32, 13, 13]
        x = F.relu(self.conv2(x))  # [B, 64, 11, 11]
        x = self.pool2(x)          # [B, 64, 5, 5]
        x = self.flatten(x)        # [B, 64*5*5]
        x = F.relu(self.fc1(x))    # [B, 128]
        x = F.softmax(self.fc2(x), dim=1)  # [B, 10]
        return x
    
model = CNNModel()

############################################
# Step 3: Convert model to ONNX
############################################

torch.onnx.export(
    model,
    torch.randn(1,1, 28,28),
    f"temp/model.onnx",
    input_names=["input"], output_names=["output"],
    dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}})

############################################
# Step 4: Generate artifacts
############################################

requires_grad = []
onnx_model = onnx.load(f"temp/model.onnx")
for node in reversed(onnx_model.graph.node):
    if node.op_type == "Conv" or node.op_type == "Gemm":
            _, second, third = node.input
            requires_grad.extend([second, third])

frozen_params = [
    param.name
    for param in onnx_model.graph.initializer
    if param.name not in requires_grad
]

artifacts.generate_artifacts(
    onnx_model,
    optimizer=artifacts.OptimType.AdamW,
    loss=artifacts.LossType.CrossEntropyLoss,
    requires_grad=requires_grad,
    frozen_params=frozen_params,
    artifact_directory="temp",
    additional_output_names=["output"])

############################################
# Step 5: Zip the files
############################################
import zipfile

files_to_zip = ["temp/checkpoint", "temp/eval_model.onnx", "temp/optimizer_model.onnx", "temp/training_model.onnx"]

with zipfile.ZipFile("artifacts.zip", "w") as zipf:
    for file in files_to_zip:
        zipf.write(file)

############################################
# Step 6: Upload the Model and Start the Benchmark.
############################################


# Log in to the dAIEdge VLab API.
api = dAIEdgeVLabAPI("setup.yaml")

# Parse the configuration file for on-device training.
# Make sure to adjust the path to your config file.
config = OnDeviceTrainingConfig("./config.json")

# Start the benchmark with the specified target, runtime, dataset and model.
id = api.startOdtBenchmark(
    TARGET,
    RUNTIME,
    MODEL,
    config
)

# Wait for the benchmark to finish and save the results.
print(f"Benchmark started with ID: {id}")

r = api.waitBenchmarkResult(id, save_path=RESULT_DIR, verbose=True)
print("Benchmark finished:", r)