準備訓練#

自定義損失#

如果需要自定義損失，使用者可以將自定義損失函式提供給 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。