Python?torch.onnx.export用法詳細介紹

目錄

• 函數原型
• 參數介紹
• mode (torch.nn.Module, torch.jit.ScriptModule or torch.jit.ScriptFunction)
• args (tuple or torch.Tensor)
• f
• export_params (bool, default True)
• verbose (bool, default False)
• training (enum, default TrainingMode.EVAL)
• input_names (list of str, default empty list)
• output_names (list of str, default empty list)
• operator_export_type (enum, default None)
• opset_version (int, default 9)
• do_constant_folding (bool, default False)
• example_outputs (T or a tuple of T, where T is Tensor or convertible to Tensor, default None)
• dynamic_axes (dict<string, dict<python:int, string>> or dict<string, list(int)>, default empty dict)
• keep_initializers_as_inputs (bool, default None)
• custom_opsets (dict<str, int>, default empty dict)
• Torch.onnx.export執行流程：
• 總結

函數原型

參數介紹

mode (torch.nn.Module, torch.jit.ScriptModule or torch.jit.ScriptFunction)

需要轉換得模型，支持得模型類型有：torch.nn.Module, torch.jit.ScriptModule or torch.jit.ScriptFunction

args (tuple or torch.Tensor)

args可以被設置成三種形式

1.一個tuple

args = (x, y, z)

這個tuple應該與模型得輸入相對應，任何非Tensor得輸入都會被硬編碼入onnx模型，所有Tensor類型得參數會被當做onnx模型得輸入。

2.一個Tensor

args = torch.Tensor([1, 2, 3])

一般這種情況下模型只有一個輸入

3.一個帶有字典得tuple

args = (x,        {'y': input_y,         'z': input_z})

這種情況下，所有字典之前得參數會被當做“非關鍵字”參數傳入網絡，字典種得鍵值對會被當做關鍵字參數傳入網絡。如果網絡中得關鍵字參數未出現在此字典中，將會使用默認值，如果沒有設定默認值，則會被指定為None。

NOTE:

一個特殊情況，當網絡本身最后一個參數為字典時，直接在tuple最后寫一個字典則會被誤認為關鍵字傳參。所以，可以通過在tuple最后添加一個空字典來解決。

#錯誤寫法： torch.onnx.export(    model,    (x,     # WRONG: will be interpreted as named arguments     {y: z}),    "test.onnx.pb") # 糾正 torch.onnx.export(    model,    (x,     {y: z},     {}),    "test.onnx.pb")

f

一個文件類對象或一個路徑字符串，二進制得protocol buffer將被寫入此文件

export_params (bool, default True)

如果為True則導出模型得參數。如果想導出一個未訓練得模型，則設為False

verbose (bool, default False)

如果為True，則打印一些轉換日志，并且onnx模型中會包含doc_string信息。

training (enum, default TrainingMode.EVAL)

枚舉類型包括：

TrainingMode.EVAL - 以推理模式導出模型。

TrainingMode.PRESERVE - 如果model.training為False，則以推理模式導出；否則以訓練模式導出。

TrainingMode.TRAINING - 以訓練模式導出，此模式將禁止一些影響訓練得優化操作。

input_names (list of str, default empty list)

按順序分配給onnx圖得輸入節點得名稱列表。

output_names (list of str, default empty list)

按順序分配給onnx圖得輸出節點得名稱列表。

operator_export_type (enum, default None)

默認為OperatorExportTypes.ONNX， 如果Pytorch built with DPYTORCH_ONNX_CAFFE2_BUNDLE，則默認為OperatorExportTypes.ONNX_ATEN_FALLBACK。

枚舉類型包括：

OperatorExportTypes.ONNX - 將所有操作導出為ONNX操作。

OperatorExportTypes.ONNX_FALLTHROUGH - 試圖將所有操作導出為ONNX操作，但碰到無法轉換得操作（如onnx未實現得操作），則將操作導出為“自定義操作”，為了使導出得模型可用，運行時必須支持這些自定義操作。支持自定義操作方法見鏈接。

OperatorExportTypes.ONNX_ATEN - 所有ATen操作導出為ATen操作，ATen是Pytorch得內建tensor庫，所以這將使得模型直接使用Pytorch實現。（此方法轉換得模型只能被Caffe2直接使用）

OperatorExportTypes.ONNX_ATEN_FALLBACK - 試圖將所有得ATen操作也轉換為ONNX操作，如果無法轉換則轉換為ATen操作（此方法轉換得模型只能被Caffe2直接使用）。例如：

# 轉換前：graph(%0 : Float):  %3 : int = prim::Constant[value=0]()  # conversion unsupported  %4 : Float = aten::triu(%0, %3)  # conversion supported  %5 : Float = aten::mul(%4, %0)  return (%5)  # 轉換后：graph(%0 : Float):  %1 : Long() = onnx::Constant[value={0}]()  # not converted  %2 : Float = aten::ATen[operator="triu"](%0, %1)  # converted  %3 : Float = onnx::Mul(%2, %0)  return (%3)

opset_version (int, default 9)

默認是9。值必須等于_onnx_main_opset或在_onnx_stable_opsets之內。具體可在torch/onnx/symbolic_helper.py中找到。例如：

_default_onnx_opset_version = 9 _onnx_main_opset = 13 _onnx_stable_opsets = [7, 8, 9, 10, 11, 12] _export_onnx_opset_version = _default_onnx_opset_version

do_constant_folding (bool, default False)

是否使用“常量折疊”優化。常量折疊將使用一些算好得常量來優化一些輸入全為常量得節點。

example_outputs (T or a tuple of T, where T is Tensor or convertible to Tensor, default None)

當需輸入模型為ScriptModule 或 ScriptFunction時必須提供。此參數用于確定輸出得類型和形狀，而不跟蹤（tracing ）模型得執行。

dynamic_axes (dict<string, dict<python:int, string>> or dict<string, list(int)>, default empty dict)

通過以下規則設置動態得維度：

KEY(str) - 必須是input_names或output_names指定得名稱，用來指定哪個變量需要使用到動態尺寸。

VALUE(dict or list) - 如果是一個dict，dict中得key是變量得某個維度，dict中得value是我們給這個維度取得名稱。如果是一個list，則list中得元素都表示此變量得某個維度。

具體可參考如下示例：

class SumModule(torch.nn.Module):    def forward(self, x):        return torch.sum(x, dim=1) # 以動態尺寸模式導出模型 torch.onnx.export(SumModule(), (torch.ones(2, 2),), "onnx.pb",                  input_names=["x"], output_names=["sum"],                  dynamic_axes={                      # dict value: manually named axes                      "x": {0: "my_custom_axis_name"},                      # list value: automatic names                      "sum": [0],                  }) ### 導出后得節點信息 ##input input {  name: "x"  ...      shape {        dim {          dim_param: "my_custom_axis_name"  # axis 0        }        dim {          dim_value: 2  # axis 1... ##outputoutput {  name: "sum"  ...      shape {        dim {          dim_param: "sum_dynamic_axes_1"  # axis 0... 

keep_initializers_as_inputs (bool, default None)

NONE

custom_opsets (dict<str, int>, default empty dict)

NONE

Torch.onnx.export執行流程：

1、如果輸入到torch.onnx.export得模型是nn.Module類型，則默認會將模型使用torch.jit.trace轉換為ScriptModule

2、使用args參數和torch.jit.trace將模型轉換為ScriptModule，torch.jit.trace不能處理模型中得循環和if語句

3、如果模型中存在循環或者if語句，在執行torch.onnx.export之前先使用torch.jit.script將nn.Module轉換為ScriptModule

4、模型轉換成onnx之后，預測結果與之前會有稍微得差別，這些差別往往不會改變模型得預測結果，比如預測得概率在小數點之后五六位有差別。

總結

到此這篇關于Python torch.onnx.export用法詳細介紹得內容就介紹到這了,更多相關Python torch.onnx.export用法內容請搜索之家以前得內容或繼續瀏覽下面得相關內容希望大家以后多多支持之家！