From a36e92d86edcdbc04c471fe4ea83b04629540b6c Mon Sep 17 00:00:00 2001 From: Baiju Meswani Date: Tue, 21 Jun 2022 17:57:44 -0700 Subject: [PATCH] Offline tooling readme (#11920) --- .../python/training/onnxblock/README.md | 76 +++++++++++++++++++ 1 file changed, 76 insertions(+) create mode 100644 orttraining/orttraining/python/training/onnxblock/README.md diff --git a/orttraining/orttraining/python/training/onnxblock/README.md b/orttraining/orttraining/python/training/onnxblock/README.md new file mode 100644 index 0000000000..75f58736f4 --- /dev/null +++ b/orttraining/orttraining/python/training/onnxblock/README.md @@ -0,0 +1,76 @@ +## Getting Started + +This is a simple guide on how to generate the training graphs using the offline tools. + +### 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 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 an 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 + +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): + # 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 model +model_path = "model.onnx" +model = onnx.load(model_path) + +# Build the training model +my_model = MyTrainingModel() +with onnxblock.onnx_model(model) as accessor: + loss_output_name = my_model("output-0", "output-1") + eval_model = accessor.eval_model + +# Save the 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 all the `onnxblocks` that are supported, please look at the [building_blocks](https://github.com/microsoft/onnxruntime/blob/training_dev/on_device_poc/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/tree/training_dev/on_device_poc/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/training_dev/on_device_poc/orttraining/orttraining/python/training/onnxblock/optim/optim.py).