1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
|
Prepare for training
=====================
Before the training can start on edge devices, the training artifacts need to be generated in an offline step.
These artifacts include:
1. The training onnx model
2. The checkpoint state
3. The optimizer onnx model
4. The eval onnx model (optional)
It is assumed that the an forward only onnx model is already available. This model can be generated by exporting the PyTorch model using the :func:`torch.onnx.export` API if using PyTorch.
.. note::
If using PyTorch to export the model, please use the following export arguments so training artifact generation can be successful:
- ``export_params``: ``True``
- ``do_constant_folding``: ``False``
- ``training``: ``torch.onnx.TrainingMode.TRAINING``
Once the forward only onnx model is available, the training artifacts can be generated using the :func:`onnxruntime.training.artifacts.generate_artifacts` API.
Sample usage:
.. code-block:: python
from onnxruntime.training import artifacts
# Load the forward only onnx model
model = onnx.load(path_to_forward_only_onnx_model)
# Generate the training artifacts
artifacts.generate_artifacts(model,
requires_grad = ["parameters", "needing", "gradients"],
frozen_params = ["parameters", "not", "needing", "gradients"],
loss = artifacts.LossType.CrossEntropyLoss,
optimizer = artifacts.OptimType.AdamW,
artifact_directory = path_to_output_artifact_directory)
.. autoclass:: onnxruntime.training.artifacts.LossType
:members:
:member-order: bysource
:undoc-members:
.. autoclass:: onnxruntime.training.artifacts.OptimType
:members:
:member-order: bysource
:undoc-members:
.. autofunction:: onnxruntime.training.artifacts.generate_artifacts
Custom Loss
++++++++++++
If a custom loss is needed, the user can provide a custom loss function to the :func:`onnxruntime.training.artifacts.generate_artifacts` API.
This is done by inheriting from the :class:`onnxruntime.training.onnxblock.Block` class and implementing the `build` method.
The following example shows how to implement a custom loss function:
Let's assume, we want to use a custom loss function with a model. For this example, we assume that our model generates
two outputs. And the custom loss function must apply a loss function on each of the outputs and perform a weighted average
on the output. Mathematically,
.. code-block:: python
loss = 0.4 * mse_loss1(output1, target1) + 0.6 * mse_loss2(output2, target2)
Since this is a custom loss function, this loss type is not exposed as an enum by `LossType` enum.
For this, we make use of `onnxblock`.
.. code-block:: python
import onnxruntime.training.onnxblock as onnxblock
from onnxruntime.training import artifacts
# Define a custom loss block that takes in two inputs
# and performs a weighted average of the losses from these
# two inputs.
class WeightedAverageLoss(onnxblock.Block):
def __init__(self):
self._loss1 = onnxblock.loss.MSELoss()
self._loss2 = onnxblock.loss.MSELoss()
self._w1 = onnxblock.blocks.Constant(0.4)
self._w2 = onnxblock.blocks.Constant(0.6)
self._add = onnxblock.blocks.Add()
self._mul = onnxblock.blocks.Mul()
def build(self, loss_input_name1, loss_input_name2):
# The build method defines how the block should be stacked on top of
# loss_input_name1 and 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"))
)
my_custom_loss = WeightedAverageLoss()
# Load the onnx model
model_path = "model.onnx"
base_model = onnx.load(model_path)
# Define the parameters that need their gradient computed
requires_grad = ["weight1", "bias1", "weight2", "bias2"]
frozen_params = ["weight3", "bias3"]
# Now, we can invoke generate_artifacts with this custom loss function
artifacts.generate_artifacts(base_model, requires_grad = requires_grad, frozen_params = frozen_params,
loss = my_custom_loss, optimizer = artifacts.OptimType.AdamW)
# Successful completion of the above call will generate 4 files in the current working directory,
# one for each of the artifacts mentioned above (training_model.onnx, eval_model.onnx, checkpoint, optimizer_model.onnx)
.. autoclass:: onnxruntime.training.onnxblock.Block
:members:
:show-inheritance:
:member-order: bysource
:inherited-members:
Advanced Usage
+++++++++++++++
`onnxblock` is a library that can be used to build complex onnx models by stacking simple blocks on top of each other. An example of this is the ability to build a custom loss function as shown above.
`onnxblock` also provides a way to build a custom forward only or training (forward + backward) onnx model through the :class:`onnxruntime.training.onnxblock.ForwardBlock` and :class:`onnxruntime.training.onnxblock.TrainingBlock` classes respectively. These blocks inherit from the base :class:`onnxruntime.training.onnxblock.Block` class and provide additional functionality to build inference and training models.
.. autoclass:: onnxruntime.training.onnxblock.ForwardBlock
:members:
:show-inheritance:
:member-order: bysource
:inherited-members:
.. autoclass:: onnxruntime.training.onnxblock.TrainingBlock
:members:
:show-inheritance:
:member-order: bysource
:inherited-members:
|
