Implementation:Microsoft Onnxruntime OrtTrainingSession

Knowledge Sources	Description
Source File	java/src/main/java/ai/onnxruntime/OrtTrainingSession.java
Repository	Microsoft/onnxruntime

Domains

• Machine Learning Training
• Model Fine-Tuning
• Checkpoint Management

Overview

OrtTrainingSession wraps an ONNX training model and provides training, evaluation, and optimization capabilities. It manages the training loop lifecycle: forward pass with gradient accumulation (trainStep), evaluation (evalStep), gradient reset (lazyResetGrad), optimizer weight updates (optimizerStep), learning rate scheduling, checkpoint save/load, and model export for inference. It also supports adding and retrieving checkpoint properties.

Description

The OrtTrainingSession class implements AutoCloseable and provides:

• Training step: Multiple trainStep() overloads compute outputs and accumulate gradients. Accepts input maps, requested outputs, pinned outputs, and run options.
• Evaluation step: Multiple evalStep() overloads compute outputs without gradient accumulation.
• Optimizer step: optimizerStep() applies gradient updates to trainable parameters.
• Gradient management: lazyResetGrad() queues gradient reset for the next training step.
• Learning rate: setLearningRate(float) and getLearningRate() for manual LR control.
• LR scheduling: registerLinearLRScheduler(warmupSteps, totalSteps, initialLR) and schedulerStep() for built-in linear warmup scheduling.
• Checkpoint management: saveCheckpoint(Path, boolean) saves state; properties can be added/retrieved via addProperty() and getFloatProperty()/getIntProperty()/getStringProperty().
• Model export: exportModelForInference(Path, String[]) exports the eval model pruned to specified output nodes.
• RNG seed: Static setSeed(long) sets the global random seed.
• Input/output names: getTrainInputNames(), getTrainOutputNames(), getEvalInputNames(), getEvalOutputNames().

The inner class OrtCheckpointState manages checkpoint loading, saving, and property storage.

Code Reference

Source Location

// File: java/src/main/java/ai/onnxruntime/OrtTrainingSession.java
// Package: ai.onnxruntime

Signature

public final class OrtTrainingSession implements AutoCloseable {
    public Set<String> getTrainInputNames();
    public Set<String> getTrainOutputNames();
    public Set<String> getEvalInputNames();
    public Set<String> getEvalOutputNames();

    public OrtSession.Result trainStep(Map<String, ? extends OnnxTensorLike> inputs) throws OrtException;
    public OrtSession.Result evalStep(Map<String, ? extends OnnxTensorLike> inputs) throws OrtException;
    public void lazyResetGrad() throws OrtException;
    public void optimizerStep() throws OrtException;
    public void optimizerStep(OrtSession.RunOptions runOptions) throws OrtException;

    public void setLearningRate(float learningRate) throws OrtException;
    public float getLearningRate() throws OrtException;
    public void registerLinearLRScheduler(long warmupSteps, long totalSteps, float initialLearningRate) throws OrtException;
    public void schedulerStep() throws OrtException;
    public static void setSeed(long seed) throws OrtException;

    public void saveCheckpoint(Path outputPath, boolean saveOptimizer) throws OrtException;
    public void addProperty(String name, float value) throws OrtException;
    public void addProperty(String name, int value) throws OrtException;
    public void addProperty(String name, String value) throws OrtException;
    public float getFloatProperty(String name) throws OrtException;
    public int getIntProperty(String name) throws OrtException;
    public String getStringProperty(String name) throws OrtException;

    public void exportModelForInference(Path outputPath, String[] outputNames) throws OrtException;
    public void close();
}

Import

import ai.onnxruntime.OrtTrainingSession;

I/O Contract

Inputs

Name	Type	Description
inputs	Map<String, ? extends OnnxTensorLike>	Named input tensors (features and targets)
learningRate	float	The learning rate to set
warmupSteps	long	Number of warmup steps for the LR scheduler
totalSteps	long	Total number of training steps
outputPath	Path	Directory path for saving checkpoints or exporting models

Outputs

Name	Type	Description
trainStep result	OrtSession.Result	Training outputs (e.g. loss)
evalStep result	OrtSession.Result	Evaluation outputs
getLearningRate()	float	Current learning rate
getFloatProperty()	float	Named float property from checkpoint

Usage Examples

import ai.onnxruntime.*;
import java.nio.file.Paths;
import java.util.*;

OrtEnvironment env = OrtEnvironment.getEnvironment();
OrtTrainingSession.setSeed(42);

try (OrtSession.SessionOptions opts = new OrtSession.SessionOptions()) {
    try (OrtTrainingSession trainingSession = env.createTrainingSession(
            "/path/to/checkpoint",
            "/path/to/training_model.onnx",
            "/path/to/eval_model.onnx",
            "/path/to/optimizer_model.onnx",
            opts)) {

        // Register a linear LR scheduler
        trainingSession.registerLinearLRScheduler(100, 1000, 0.001f);

        // Training loop
        for (int epoch = 0; epoch < 10; epoch++) {
            trainingSession.lazyResetGrad();

            // Prepare inputs (features + labels)
            Map<String, OnnxTensorLike> inputs = new HashMap<>();
            // ... populate inputs ...

            try (OrtSession.Result result = trainingSession.trainStep(inputs)) {
                float loss = ((float[][]) result.get(0).getValue())[0][0];
                System.out.println("Epoch " + epoch + " loss: " + loss);
            }

            trainingSession.optimizerStep();
            trainingSession.schedulerStep();
        }

        // Save checkpoint
        trainingSession.addProperty("epoch", 10);
        trainingSession.saveCheckpoint(Paths.get("/path/to/output_checkpoint"), true);

        // Export for inference
        trainingSession.exportModelForInference(
            Paths.get("/path/to/inference_model.onnx"),
            new String[]{"output"});
    }
}

Related Pages

• OrtEnvironment.java - Factory for creating training sessions
• OrtSession.java - Inference session and shared inner classes
• ai_onnxruntime_OrtTrainingSession.c - JNI native implementation
• OnnxTensor.java - Tensor inputs for training and evaluation