View source for Implementation:Microsoft Onnxruntime OrtTrainingSession

{{PageInfo|type=Implementation|title=Microsoft_Onnxruntime_OrtTrainingSession}}
{{PageInfo|type=Implementation|title=Microsoft_Onnxruntime_OrtTrainingSession}}
{| class="wikitable"
! 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 (<code>trainStep</code>), evaluation (<code>evalStep</code>), gradient reset (<code>lazyResetGrad</code>), optimizer weight updates (<code>optimizerStep</code>), learning rate scheduling, checkpoint save/load, and model export for inference. It also supports adding and retrieving checkpoint properties.

== Description ==
The <code>OrtTrainingSession</code> class implements <code>AutoCloseable</code> and provides:

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

The inner class <code>OrtCheckpointState</code> manages checkpoint loading, saving, and property storage.

== Code Reference ==

=== Source Location ===
<syntaxhighlight lang="java">
// File: java/src/main/java/ai/onnxruntime/OrtTrainingSession.java
// Package: ai.onnxruntime
</syntaxhighlight>

=== Signature ===
<syntaxhighlight lang="java">
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();
}
</syntaxhighlight>

=== Import ===
<syntaxhighlight lang="java">
import ai.onnxruntime.OrtTrainingSession;
</syntaxhighlight>

== I/O Contract ==

=== Inputs ===
{| class="wikitable"
! 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 ===
{| class="wikitable"
! 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 ==
<syntaxhighlight lang="java">
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"});
    }
}
</syntaxhighlight>

== Related Pages ==
* [[Microsoft_Onnxruntime_OrtEnvironment|OrtEnvironment.java]] - Factory for creating training sessions
* [[Microsoft_Onnxruntime_OrtSession|OrtSession.java]] - Inference session and shared inner classes
* [[Microsoft_Onnxruntime_OrtTrainingSession_JNI|ai_onnxruntime_OrtTrainingSession.c]] - JNI native implementation
* [[Microsoft_Onnxruntime_OnnxTensor|OnnxTensor.java]] - Tensor inputs for training and evaluation

[[Category:Implementations]]

[[Category:Implementations]]