准备训练#

在边缘设备上开始训练之前，需要在离线步骤中生成训练工件。

这些工件包括

训练 ONNX 模型
检查点状态
优化器 ONNX 模型
评估 ONNX 模型（可选）

假设已有一个仅包含前向计算的 ONNX 模型。如果使用 PyTorch，可以通过 torch.onnx.export() API 导出 PyTorch 模型来生成此模型。

注意

如果使用 PyTorch 导出模型，请使用以下导出参数，以确保训练工件生成成功

export_params: True
do_constant_folding: False
training: torch.onnx.TrainingMode.TRAINING

一旦仅包含前向计算的 ONNX 模型可用，就可以使用 onnxruntime.training.artifacts.generate_artifacts() API 生成训练工件。

示例用法

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)

class onnxruntime.training.artifacts.LossType(value)[source]#

添加到训练模型的损失类型。

与 generate_artifacts 函数的 loss 参数一起使用。

MSELoss = 1#

CrossEntropyLoss = 2#

BCEWithLogitsLoss = 3#

L1Loss = 4#

class onnxruntime.training.artifacts.OptimType(value)[source]#

生成训练优化器模型时使用的优化器类型。

与 generate_artifacts 函数的 optimizer 参数一起使用。

AdamW = 1#

SGD = 2#

onnxruntime.training.artifacts.generate_artifacts(model: Union[ModelProto, str], requires_grad: Optional[List[str]] = None, frozen_params: Optional[List[str]] = None, loss: Optional[Union[LossType, Block]] = None, optimizer: Optional[Union[OptimType, Block]] = None, artifact_directory: Optional[Union[str, bytes, PathLike]] = None, prefix: str = '', ort_format: bool =False, custom_op_library: Optional[Union[str, bytes, PathLike]] =None, additional_output_names: Optional[List[str]] =None, nominal_checkpoint: bool =False, loss_input_names: Optional[List[str]] =None) → None[source]#

生成使用 ORT 训练 API 进行训练所需的工件。

此函数生成以下工件

训练模型 (onnx.ModelProto)：包含基础模型图、损失子图和梯度图。
评估模型 (onnx.ModelProto)：包含基础模型图和损失子图
检查点（目录）：包含模型参数。
优化器模型 (onnx.ModelProto)：包含优化器图的模型。

所有生成的 ModelProto 将使用由 model 定义的相同 opsets。

参数:

model – 用于生成梯度图的基础模型或基础模型的路径。对于大于 2GB 的模型，请使用基础模型的路径。
requires_grad – 需要计算梯度的模型参数名称列表
frozen_params – 应冻结的模型参数名称列表。
loss – 用于训练的损失函数枚举或 onnxblock。如果为 None，则图中不添加损失节点。
optimizer – 用于训练的优化器枚举或 onnxblock。如果为 None，则不生成优化器模型。
artifact_directory – 保存生成工件的目录。如果为 None，则使用当前工作目录。
prefix – 用于生成工件的前缀。如果未指定，则不使用前缀。
ort_format – 是否以 ORT 格式保存生成工件。默认为 False。
custom_op_library – 自定义操作库的路径。如果未指定，则不使用自定义操作库。
additional_output_names – 除了损失输出之外，添加到训练/评估模型的附加输出名称列表。默认为 None。
nominal_checkpoint – 除了完整检查点外，是否生成名义检查点。默认为 False。名义检查点是一个包含有关模型参数名义信息的检查点。它可用于在设备上构建训练模型时减少开销，以及减小设备上应用程序打包的检查点大小。
loss_input_names – 指定专门用于损失计算的输入名称列表。提供时，只有这些输入将传递给损失函数。如果为 None，则所有图输出都将传递给损失函数。

引发:

RuntimeError – 如果提供的损失既不是支持的损失之一，也不是 onnxblock.Block 的实例
RuntimeError – 如果提供的优化器不是支持的优化器之一。

自定义损失#

如果需要自定义损失，用户可以将自定义损失函数提供给 onnxruntime.training.artifacts.generate_artifacts() API。这可以通过继承 onnxruntime.training.onnxblock.Block 类并实现 build 方法来完成。

以下示例展示了如何实现自定义损失函数

假设我们要将自定义损失函数与模型一起使用。在此示例中，我们假设模型生成两个输出。自定义损失函数必须对每个输出应用损失函数，并对输出执行加权平均。数学上，

loss = 0.4 * mse_loss1(output1, target1) + 0.6 * mse_loss2(output2, target2)

由于这是自定义损失函数，因此此损失类型不作为 LossType 枚举公开。

为此，我们使用 onnxblock。

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)

class onnxruntime.training.onnxblock.Block(temp_file_name='temp.onnx')[source]#

基类：ABC

所有可以相互堆叠的构建块的基类。

所有需要操作模型的块都必须继承此类。子类对 build 方法的实现必须返回块的中间输出名称。

子类对 build 方法的实现必须根据需要操作基础模型，但操作后的模型必须是有效的（由 onnx 检查器判定）。

base#

子类可以操作的基础模型。

类型:: onnx.ModelProto

abstract build(*args, **kwargs)[source]#

通过在此函数的输入之上堆叠块来定制模型。

此方法必须由子类重写。

infer_shapes_on_base()[source]#

对全局模型执行形状推断。如果使用了路径，则使用 infer_shapes_path API 支持具有外部数据的模型。

返回形状推断后的 ModelProto。

高级用法#

onnxblock 是一个库，可以通过相互堆叠简单的块来构建复杂的 ONNX 模型。例如，上面所示的构建自定义损失函数的功能。

onnxblock 还提供了一种方法，通过 onnxruntime.training.onnxblock.ForwardBlock 和 onnxruntime.training.onnxblock.TrainingBlock 类分别构建自定义的仅前向或训练（前向+后向）ONNX 模型。这些块继承自基础 onnxruntime.training.onnxblock.Block 类，并提供构建推理和训练模型的附加功能。

class onnxruntime.training.onnxblock.ForwardBlock[source]#

基类：Block

所有需要自动构建前向模型的块的基类。

需要通过在现有模型之上堆叠块来构建前向模型的块必须继承此类。子类对 build 方法的实现必须返回图输出的名称。此块将自动将该输出注册为图输出并构建模型。

示例

>>> class MyForwardBlock(ForwardBlock):
>>>     def __init__(self):
>>>         super().__init__()
>>>         self.loss = onnxblock.loss.CrossEntropyLoss()
>>>
>>>     def build(self, loss_input_name: str):
>>>         # Add a cross entropy loss on top of the output so far (loss_input_name)
>>>         return self.loss(loss_input_name)

以上示例将自动构建前向图，该图由现有模型和堆叠在其上的交叉熵损失函数组成。

abstract build(*args, **kwargs)[source]#

通过在此函数的输入之上堆叠块来定制此模型的前向图。

此方法应由子类重写。此方法的输出应为图输出的名称。

to_model_proto()[source]#

返回前向模型。

返回:: 前向模型。
返回类型:: 模型 (onnx.ModelProto)
引发:: RuntimeError – 如果尚未调用 build 方法（即尚未构建前向模型）。

infer_shapes_on_base()#

对全局模型执行形状推断。如果使用了路径，则使用 infer_shapes_path API 支持具有外部数据的模型。

返回形状推断后的 ModelProto。

class onnxruntime.training.onnxblock.TrainingBlock[source]#

基类：Block

所有需要自动构建梯度模型的块的基类。

需要根据块的输出来计算梯度图的块必须继承此类。子类对 build 方法的实现必须返回反向传播开始的输出名称（通常是损失函数的输出名称）。

示例

>>> class MyTrainingBlock(TrainingBlock):
>>>     def __init__(self):
>>>         super().__init__()
>>>         self.loss = onnxblock.loss.CrossEntropyLoss()
>>>
>>>     def build(self, loss_input_name: str):
>>>         # Add a cross entropy loss on top of the output so far (loss_input_name)
>>>         return self.loss(loss_input_name)

以上示例将自动构建从损失函数输出开始的整个模型的梯度图。

abstract build(*args, **kwargs)[source]#

通过在此函数的输入之上堆叠块来定制此模型的前向图。

此方法应由子类重写。此方法的输出应为反向传播开始的输出名称（通常是损失函数的输出名称）。

requires_grad(argument_name: str, value: bool =True)[source]#

指定参数是否需要梯度。

自动微分将仅计算需要梯度的参数的梯度图。默认情况下，没有参数需要梯度。用户必须明确指定哪些参数需要梯度。

参数:

argument_name (str) – 需要/不需要梯度的参数名称。
value (bool) – 如果参数需要梯度，则为 True，否则为 False。

parameters() → Tuple[List[TensorProto], List[TensorProto]][source]#

模型的可训练参数和不可训练参数（冻结参数）。

构建训练模型时提取的模型参数将由此方法返回。

请注意，在构建训练模型之前参数是未知的。因此，如果在构建训练模型之前调用此方法，将引发异常。

返回:: 模型的可训练参数。frozen_params (onnx.TensorProto 列表)：模型的不可训练参数。
返回类型:: trainable_params (list of onnx.TensorProto)
引发:: RuntimeError – 如果尚未调用 build 方法（即尚未构建训练模型）。

to_model_proto() → Tuple[ModelProto, ModelProto][source]#

返回训练和评估模型。

一旦构建了梯度图，就可以通过调用此方法检索训练和评估模型。

返回:: 训练模型。eval_model (onnx.ModelProto): 评估模型。
返回类型:: training_model (onnx.ModelProto)
引发:: RuntimeError – 如果尚未调用 build 方法（即尚未构建训练模型）。

infer_shapes_on_base()#

对全局模型执行形状推断。如果使用了路径，则使用 infer_shapes_path API 支持具有外部数据的模型。

返回形状推断后的 ModelProto。