## Getting Started

This is a simple guide on how to generate the training graphs using the offline tools.

### What's needed for training?
The on-device training APIs need the following files for performing training
1. The training onnx model.
2. The eval onnx model (optional).
3. The optimizer onnx model.
4. The checkpoint file.

The offline tooling helps generating the above files. Read along to know how.

### The Training and Eval Graph

The training graph is generated in place using the offline tools. Let's assume that an inference only onnx model has already been generated. Here, we require that the user's onnx model is generated with the parameters embedded inside the exported model, i.e. with the parameters `export_params=True` and `training=torch.onnx.TrainingMode.TRAINING`. Let's also assume you're working with a model that has two outputs. Let's call them `output-0` and `output-1`. The training graph can be generated by writing the following code. The eval model can be obtained as a by product.

```py
import onnxruntime.training.onnxblock as onnxblock

# Define how the training model should be constructed.
# In this example, the training model is constructed by stacking MSELoss
# blocks on top of the inference model (with two outputs) and taking a
# weighted average of the two losses to compute the final loss.
# And the onnxblock tooling automatically takes care of building the
# training graph for this model.
class MyTrainingModel(onnxblock.TrainingModel):
    def __init__(self):
        super(MyTrainingModel, self).__init__()
        self._loss1 = onnxblock.loss.MSELoss()
        self._loss2 = onnxblock.loss.MSELoss()
        self._w1 = onnxblock.building_blocks.Constant(0.6)
        self._w2 = onnxblock.building_blocks.Constant(0.4)
        self._add = onnxblock.building_blocks.Add()
        self._mul = onnxblock.building_blocks.Mul()

    def build(self, loss_input_name1, loss_input_name2):
        # The build method defines how the forward part of the training
        # graph should be built and must be defined by the user
        # inheriting from onnxblock.TrainingModel.

        # Returns weighted average of the two losses
        return self._add(
            self._mul(self._w1(), self._loss1(loss_input_name1, target_name="target1")),
            self._mul(self._w2(), self._loss2(loss_input_name2, target_name="target2"))
        )

# Load the onnx model
model_path = "model.onnx"
model = onnx.load(model_path)

# Build the training model by adding the loss block and building the
# gradient graph. This is achieved by invoking an instance of
# the onnxblock.TrainingModel.
my_model = MyTrainingModel()
with onnxblock.onnx_model(model) as accessor:
    # Invoking the model builder will add the loss block on top
    # of the loaded onnx model "model" inplace.
    loss_output_name = my_model("output-0", "output-1")
    # The training model is built in place on top of the loaded onnx model "model".
    # But if needed, it can also be accessed by accessor.model.
    # Eval model is obtained as a by-product of building
    # the training model.
    eval_model = accessor.eval_model

# Save the generated onnx models
output_model_path = "training_model.onnx"
output_eval_model_path = "eval_model.onnx"
onnx.save(model, output_model_path)
onnx.save(eval_model, output_eval_model_path)
```

### The Optimizer Graph

The optimizer graph can be generated in a similar way. An example snippet is as follows:

```py
cli_grad_norm = onnxblock.optim.ClipGradNorm(2.5)
optimizer = onnxblock.optim.AdamW(clip_grad=clip_grad_norm)
optimizer_model = None
with onnxblock.onnx_model() as accessor:
    optimizer_outputs = optimizer(my_model.parameters())
    optimizer_model = accessor.model

# save the optimizer model
output_optimizer_model_path = "adamw_optimizer.onnx"
onnx.save(optimizer_model, output_optimizer_model_path)
```

### The Checkpoint

The model checkpoint can be saved using instance of the `TrainingModel`'s `parameters()` method.

```py
# save checkpoint
output_checkpoint_path = "checkpoint.ckpt"
onnxblock.save_checkpoint(my_model.parameters(), output_checkpoint_path)
```

Once the models and checkpoint have been generated, they can be loaded in the online training step and executed.
For an example on how the online training loop should be written given these generated files, refer to this
[sample trainer](https://github.com/microsoft/onnxruntime/blob/main/orttraining/orttraining/test/training_api/trainer/trainer.cc).
For all the `onnxblocks` that are supported, please look at the [building_blocks](https://github.com/microsoft/onnxruntime/blob/main/orttraining/orttraining/python/training/onnxblock/building_blocks.py).
For all the loss blocks that are supported, please look at the [loss blocks](https://github.com/microsoft/onnxruntime/blob/main/orttraining/orttraining/python/training/onnxblock/loss/loss.py).
For all the optimizer blocks that are supported, please look at the [optim blocks](https://github.com/microsoft/onnxruntime/blob/main/orttraining/orttraining/python/training/onnxblock/optim/optim.py).
