torch.onnx.export(model, args, f, export_params=True, verbose=False, training=False, input_names=None, output_names=None, aten=False, export_raw_ir=False, operator_export_type=None, opset_version=None, _retain_param_name=True, do_constant_folding=False, example_outputs=None, strip_doc_string=True, dynamic_axes=None, keep_initializers_as_inputs=None)
功能:
將模型以ONNX格式導出並保存.
參數:
- model (torch.nn.Module) – 要導出的模型.
- args (tuple of arguments) – 模型的輸入, 任何非Tensor參數都將硬編碼到導出的模型中；任何Tensor參數都將成爲導出的模型的輸入，並按照他們在args中出現的順序輸入。因爲export運行模型，所以我們需要提供一個輸入張量x。只要是正確的類型和大小，其中的值就可以是隨機的。請注意，除非指定爲動態軸，否則輸入尺寸將在導出的ONNX圖形中固定爲所有輸入尺寸。在此示例中，我們使用輸入batch_size 1導出模型，但隨後dynamic_axes 在torch.onnx.export()。因此，導出的模型將接受大小爲[batch_size，3、100、100]的輸入，其中batch_size可以是可變的。
- export_params (bool, default True) – 如果指定爲True或默認, 參數也會被導出. 如果你要導出一個沒訓練過的就設爲 False.
- verbose (bool, default False) - 如果指定，我們將打印出一個導出軌跡的調試描述。
- training (bool, default False) - 在訓練模式下導出模型。目前，ONNX導出的模型只是爲了做推斷，所以你通常不需要將其設置爲True。
- input_names (list of strings, default empty list) – 按順序分配名稱到圖中的輸入節點
- output_names (list of strings, default empty list) –按順序分配名稱到圖中的輸出節點
- dynamic_axes – {‘input’ : {0 : ‘batch_size’}, ‘output’ : {0 : ‘batch_size’}}) # variable lenght axes
舉例:
shape(input_1) = ('b', 3, 'w', 'h')
and shape(input_2) = ('b', 4)
and shape(output) = ('b', 'd', 5)
torch.onnx.export(model, args, f, export_params=True, verbose=False, training=False, input_names=None, output_names=None, dynamic_axes = {‘input_1’:[0, 2, 3], ‘input_2’:[0], ‘output’:[0, 1]})