Offline tooling readme (#11920)

This commit is contained in:
Baiju Meswani 2022-06-21 17:57:44 -07:00 committed by GitHub
parent a3ec2d6d15
commit a36e92d86e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -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).