PyTorch torch.onnx

2020-09-15 10:47 更新

原文: PyTorch torch.onnx

示例:從 PyTorch 到 ONNX 的端到端 AlexNet

這是一個簡單的腳本,可以將 Torchvision 中定義的經(jīng)過預(yù)訓練的 AlexNet 導出到 ONNX 中。 它運行一輪推斷,然后將生成的跟蹤模型保存到alexnet.onnx

import torch
import torchvision


dummy_input = torch.randn(10, 3, 224, 224, device='cuda')
model = torchvision.models.alexnet(pretrained=True).cuda()


## Providing input and output names sets the display names for values
## within the model's graph. Setting these does not change the semantics
## of the graph; it is only for readability.
## ## The inputs to the network consist of the flat list of inputs (i.e.
## the values you would pass to the forward() method) followed by the
## flat list of parameters. You can partially specify names, i.e. provide
## a list here shorter than the number of inputs to the model, and we will
## only set that subset of names, starting from the beginning.
input_names = [ "actual_input_1" ] + [ "learned_%d" % i for i in range(16) ]
output_names = [ "output1" ]


torch.onnx.export(model, dummy_input, "alexnet.onnx", verbose=True, input_names=input_names, output_names=output_names)

生成的alexnet.onnx是二進制 protobuf 文件,其中包含您導出的模型的網(wǎng)絡(luò)結(jié)構(gòu)和參數(shù)(在本例中為 AlexNet)。 關(guān)鍵字參數(shù)verbose=True使導出程序打印出人類可讀的網(wǎng)絡(luò)表示形式:

## These are the inputs and parameters to the network, which have taken on
## the names we specified earlier.
graph(%actual_input_1 : Float(10, 3, 224, 224)
      %learned_0 : Float(64, 3, 11, 11)
      %learned_1 : Float(64)
      %learned_2 : Float(192, 64, 5, 5)
      %learned_3 : Float(192)
      # ---- omitted for brevity ----
      %learned_14 : Float(1000, 4096)
      %learned_15 : Float(1000)) {
  # Every statement consists of some output tensors (and their types),
  # the operator to be run (with its attributes, e.g., kernels, strides,
  # etc.), its input tensors (%actual_input_1, %learned_0, %learned_1)
  %17 : Float(10, 64, 55, 55) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[11, 11], pads=[2, 2, 2, 2], strides=[4, 4]](%actual_input_1, %learned_0, %learned_1), scope: AlexNet/Sequential[features]/Conv2d[0]
  %18 : Float(10, 64, 55, 55) = onnx::Relu(%17), scope: AlexNet/Sequential[features]/ReLU[1]
  %19 : Float(10, 64, 27, 27) = onnx::MaxPool[kernel_shape=[3, 3], pads=[0, 0, 0, 0], strides=[2, 2]](%18), scope: AlexNet/Sequential[features]/MaxPool2d[2]
  # ---- omitted for brevity ----
  %29 : Float(10, 256, 6, 6) = onnx::MaxPool[kernel_shape=[3, 3], pads=[0, 0, 0, 0], strides=[2, 2]](%28), scope: AlexNet/Sequential[features]/MaxPool2d[12]
  # Dynamic means that the shape is not known. This may be because of a
  # limitation of our implementation (which we would like to fix in a
  # future release) or shapes which are truly dynamic.
  %30 : Dynamic = onnx::Shape(%29), scope: AlexNet
  %31 : Dynamic = onnx::Slice[axes=[0], ends=[1], starts=[0]](%30), scope: AlexNet
  %32 : Long() = onnx::Squeeze[axes=[0]](%31), scope: AlexNet
  %33 : Long() = onnx::Constant[value={9216}](), scope: AlexNet
  # ---- omitted for brevity ----
  %output1 : Float(10, 1000) = onnx::Gemm[alpha=1, beta=1, broadcast=1, transB=1](%45, %learned_14, %learned_15), scope: AlexNet/Sequential[classifier]/Linear[6]
  return (%output1);
}

您還可以使用 ONNX 庫來驗證 protobuf。 您可以使用 conda 安裝ONNX

conda install -c conda-forge onnx

然后,您可以運行:

import onnx


## Load the ONNX model
model = onnx.load("alexnet.onnx")


## Check that the IR is well formed
onnx.checker.check_model(model)


## Print a human readable representation of the graph
onnx.helper.printable_graph(model.graph)

要使用 caffe2 運行導出的腳本,您將需要安裝 <cite>caffe2</cite> :如果尚未安裝,請按照安裝說明進行操作。

一旦安裝了這些,就可以將后端用于 Caffe2:

## ...continuing from above
import caffe2.python.onnx.backend as backend
import numpy as np


rep = backend.prepare(model, device="CUDA:0") # or "CPU"
## For the Caffe2 backend:
##     rep.predict_net is the Caffe2 protobuf for the network
##     rep.workspace is the Caffe2 workspace for the network
##       (see the class caffe2.python.onnx.backend.Workspace)
outputs = rep.run(np.random.randn(10, 3, 224, 224).astype(np.float32))
## To run networks with more than one input, pass a tuple
## rather than a single numpy ndarray.
print(outputs[0])

您還可以使用 ONNX Runtime 運行導出的模型,您將需要安裝 <cite>ONNX Runtime</cite> :請按照這些說明進行操作。

一旦安裝了這些,就可以將后端用于 ONNX Runtime:

## ...continuing from above
import onnxruntime as ort


ort_session = ort.InferenceSession('alexnet.onnx')


outputs = ort_session.run(None, {'actual_input_1': np.random.randn(10, 3, 224, 224).astype(np.float32)})


print(outputs[0])

這是將 SuperResolution 模型導出到 ONNX 的另一本教程。 。

將來,其他框架也會有后端。

跟蹤與腳本編寫

ONNX 導出器可以是基于跟蹤的和基于腳本的導出器。

  • 基于跟蹤的表示它通過執(zhí)行一次模型并導出在此運行期間實際運行的運算符進行操作。 這意味著如果您的模型是動態(tài)的,例如根據(jù)輸入數(shù)據(jù)更改行為,則導出將不準確。 同樣,跟蹤可能僅對特定的輸入大小才有效(這是我們在跟蹤時需要顯式輸入的原因之一。)我們建議檢查模型跟蹤并確保所跟蹤的運算符看起來合理。 如果您的模型包含控制循環(huán)(如 for 循環(huán))和 if 條件,則基于基于跟蹤的導出器將展開循環(huán)以及 if 條件,并導出與此運行完全相同的靜態(tài)圖形。 如果要使用動態(tài)控制流導出模型,則需要使用基于腳本的導出器。
  • 基于腳本的表示您要導出的模型是 ScriptModule 。 <cite>ScriptModule</cite> 是 <cite>TorchScript</cite> 中的核心數(shù)據(jù)結(jié)構(gòu), <cite>TorchScript</cite> 是 Python 語言的子集,可從 PyTorch 代碼創(chuàng)建可序列化和可優(yōu)化的模型。

我們允許混合跟蹤和腳本編寫。 您可以組合跟蹤和腳本以適合模型部分的特定要求。 看看這個例子:

import torch


## Trace-based only


class LoopModel(torch.nn.Module):
    def forward(self, x, y):
        for i in range(y):
            x = x + i
        return x


model = LoopModel()
dummy_input = torch.ones(2, 3, dtype=torch.long)
loop_count = torch.tensor(5, dtype=torch.long)


torch.onnx.export(model, (dummy_input, loop_count), 'loop.onnx', verbose=True)

使用基于跟蹤的導出器,我們得到結(jié)果 ONNX 圖,該圖展開了 for 循環(huán):

graph(%0 : Long(2, 3),
      %1 : Long()):
  %2 : Tensor = onnx::Constant[value={1}]()
  %3 : Tensor = onnx::Add(%0, %2)
  %4 : Tensor = onnx::Constant[value={2}]()
  %5 : Tensor = onnx::Add(%3, %4)
  %6 : Tensor = onnx::Constant[value={3}]()
  %7 : Tensor = onnx::Add(%5, %6)
  %8 : Tensor = onnx::Constant[value={4}]()
  %9 : Tensor = onnx::Add(%7, %8)
  return (%9)

為了利用基于腳本的導出器捕獲動態(tài)循環(huán),我們可以在腳本中編寫循環(huán),然后從常規(guī) nn.Module 中調(diào)用它:

## Mixing tracing and scripting


@torch.jit.script
def loop(x, y):
    for i in range(int(y)):
        x = x + i
    return x


class LoopModel2(torch.nn.Module):
    def forward(self, x, y):
        return loop(x, y)


model = LoopModel2()
dummy_input = torch.ones(2, 3, dtype=torch.long)
loop_count = torch.tensor(5, dtype=torch.long)
torch.onnx.export(model, (dummy_input, loop_count), 'loop.onnx', verbose=True,
                  input_names=['input_data', 'loop_range'])

現(xiàn)在,導出的 ONNX 圖變?yōu)椋?/p>

graph(%input_data : Long(2, 3),
      %loop_range : Long()):
  %2 : Long() = onnx::Constant[value={1}](), scope: LoopModel2/loop
  %3 : Tensor = onnx::Cast[to=9](%2)
  %4 : Long(2, 3) = onnx::Loop(%loop_range, %3, %input_data), scope: LoopModel2/loop # custom_loop.py:240:5
    block0(%i.1 : Long(), %cond : bool, %x.6 : Long(2, 3)):
      %8 : Long(2, 3) = onnx::Add(%x.6, %i.1), scope: LoopModel2/loop # custom_loop.py:241:13
      %9 : Tensor = onnx::Cast[to=9](%2)
      -> (%9, %8)
  return (%4)

動態(tài)控制流已正確捕獲。 我們可以在具有不同循環(huán)范圍的后端進行驗證。

import caffe2.python.onnx.backend as backend
import numpy as np
import onnx
model = onnx.load('loop.onnx')


rep = backend.prepare(model)
outputs = rep.run((dummy_input.numpy(), np.array(9).astype(np.int64)))
print(outputs[0])
#[[37 37 37]
## [37 37 37]]


import onnxruntime as ort
ort_sess = ort.InferenceSession('loop.onnx')
outputs = ort_sess.run(None, {'input_data': dummy_input.numpy(),
                              'loop_range': np.array(9).astype(np.int64)})
print(outputs)
#[array([[37, 37, 37],
##       [37, 37, 37]], dtype=int64)]

局限性

  • 導出中目前不支持張量就地索引分配,例如 <cite>data [index] = new_data</cite> 。 解決此類問題的一種方法是使用運算符<cite>散布</cite>,顯式更新原始張量。

  data = torch.zeros(3, 4)
  index = torch.tensor(1)
  new_data = torch.arange(4).to(torch.float32)

  
  # Assigning to left hand side indexing is not supported in exporting.
  # class InPlaceIndexedAssignment(torch.nn.Module):
  # def forward(self, data, index, new_data):
  #     data[index] = new_data
  #     return data

  
  class InPlaceIndexedAssignmentONNX(torch.nn.Module):
      def forward(self, data, index, new_data):
          new_data = new_data.unsqueeze(0)
          index = index.expand(1, new_data.size(1))
          data.scatter_(0, index, new_data)
          return data

  
  out = InPlaceIndexedAssignmentONNX()(data, index, new_data)

  
  torch.onnx.export(InPlaceIndexedAssignmentONNX(), (data, index, new_data), 'inplace_assign.onnx')

  
  # caffe2
  import caffe2.python.onnx.backend as backend
  import onnx

  
  onnx_model = onnx.load('inplace_assign.onnx')
  rep = backend.prepare(onnx_model)
  out_caffe2 = rep.run((torch.zeros(3, 4).numpy(), index.numpy(), new_data.numpy()))

  
  assert torch.all(torch.eq(out, torch.tensor(out_caffe2)))

  
  # onnxruntime
  import onnxruntime
  sess = onnxruntime.InferenceSession('inplace_assign.onnx')
  out_ort = sess.run(None, {
      sess.get_inputs()[0].name: torch.zeros(3, 4).numpy(),
      sess.get_inputs()[1].name: index.numpy(),
      sess.get_inputs()[2].name: new_data.numpy(),
  })

  
  assert torch.all(torch.eq(out, torch.tensor(out_ort)))

  • ONNX 中沒有張量列表的概念。 沒有這個概念,很難導出消耗或產(chǎn)生張量列表的運算符,尤其是在導出時不知道張量列表的長度的情況下。

  x = torch.tensor([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]])

  
  # This is not exportable
  class Model(torch.nn.Module):
      def forward(self, x):
          return x.unbind(0)

  
  # This is exportable.
  # Note that in this example we know the split operator will always produce exactly three outputs,
  # Thus we can export to ONNX without using tensor list.
  class AnotherModel(torch.nn.Module):
      def forward(self, x):
          return [torch.squeeze(out, 0) for out in torch.split(x, [1,1,1], dim=0)]

  • 僅將元組,列表和變量作為 JIT 輸入/輸出支持。 也接受字典和字符串,但不建議使用它們。 用戶需要仔細驗證自己的字典輸入,并記住動態(tài)查詢不可用。

  • PyTorch 和 ONNX 后端(Caffe2,ONNX 運行時等)通常具有一些數(shù)字差異的運算符實現(xiàn)。 根據(jù)模型結(jié)構(gòu)的不同,這些差異可能可以忽略不計,但是它們也可能導致行為上的重大差異(尤其是在未經(jīng)訓練的模型上。)我們允許 Caffe2 直接調(diào)用運算符的 Torch 實現(xiàn),以在精度很重要時幫助您消除這些差異。 ,并記錄這些差異。

支持的運營商

支持以下運算符:

  • 批量標準
  • ConstantPadNd
  • 轉(zhuǎn)換
  • 退出
  • 嵌入(不支持可選參數(shù))
  • FeatureDropout(不支持訓練模式)
  • 指數(shù)
  • MaxPool1d
  • MaxPool2d
  • MaxPool3d
  • RNN
  • 腹肌
  • 阿科斯
  • adaptive_avg_pool1d
  • adaptive_avg_pool2d
  • adaptive_avg_pool3d
  • adaptive_max_pool1d
  • adaptive_max_pool2d
  • adaptive_max_pool3d
  • 添加(不支持非零 Alpha)
  • addmm
  • 范圍
  • argmax
  • 精氨酸
  • 阿辛
  • 曬黑
  • avg_pool1d
  • avg_pool2d
  • avg_pool2d
  • avg_pool3d
  • Baddbmm
  • 細胞
  • 最大鉗位
  • 最小鉗位
  • 康卡特
  • cos
  • cumsum
  • 暗淡的
  • div
  • 退出
  • lu
  • 空的
  • 空的喜歡
  • 當量
  • 埃爾夫
  • 經(jīng)驗值
  • 擴大
  • expand_as
  • 展平
  • 地板
  • frobenius_norm
  • 充分
  • 滿喜歡
  • 收集
  • ge
  • 格魯
  • 谷氨酸
  • gt
  • Hardtanh
  • index_
  • index_fill
  • index_select
  • instance_norm
  • 伊斯南
  • layer_norm
  • leaky_relu
  • 日志
  • log1p
  • 日志 2
  • log_sigmoid
  • log_softmax
  • 對數(shù)表達式
  • lt
  • masked_fill
  • 最大值
  • 意思
  • 毫米
  • 多項式
  • 狹窄
  • NE
  • 負數(shù)
  • 非零
  • 規(guī)范
  • 那些
  • 喜歡
  • 要么
  • 置換
  • pixel_shuffle
  • 戰(zhàn)俘
  • prelu(不支持輸入通道之間共享的單個權(quán)重)
  • 產(chǎn)品
  • 蘭德
  • 蘭德
  • randn_like
  • 倒數(shù)
  • Reflection_pad
  • 露露
  • 重復(fù)
  • 復(fù)制墊
  • 重塑
  • reshape_as
  • 回合
  • 雷雷魯
  • rsqrt
  • 訂閱
  • 分散
  • scatter_add
  • 選擇
  • 塞盧
  • 乙狀結(jié)腸
  • 標志
  • 尺寸
  • 切片
  • 軟最大
  • 軟加
  • 分類
  • 分裂
  • sqrt
  • 性病
  • 子(不支持非零 Alpha)
  • ?
  • 棕褐色
  • 閾值(不支持非零閾值/非零值)
  • 托普
  • 轉(zhuǎn)置
  • type_as
  • 展開(與 ATen-Caffe2 集成的實驗支持)
  • 獨特
  • 松開
  • upsample_nearest1d
  • upsample_nearest2d
  • upsample_nearest3d
  • 視圖
  • 哪里
  • zeros_like

上面設(shè)置的運算符足以導出以下模型:

  • 亞歷克斯網(wǎng)
  • DCGAN
  • 密集網(wǎng)
  • 初始階段(警告:此模型對操作員實施的更改高度敏感)
  • ResNet
  • 超分辨率
  • VGG
  • word_language_model

添加對運營商的支持

為操作員添加導出支持是的高級用法。 為此,開發(fā)人員需要觸摸 PyTorch 的源代碼。 請按照說明從源代碼安裝 PyTorch。 如果所需的運算符在 ONNX 中已標準化,則應(yīng)該容易添加對導出此類運算符的支持(為該運算符添加符號功能)。 要確認操作員是否標準化,請檢查 ONNX 操作員列表

ATen 運算符

如果該運算符是 ATen 運算符,則意味著您可以在torch/csrc/autograd/generated/VariableType.h中找到該函數(shù)的聲明(可在 PyTorch 安裝目錄的生成代碼中找到),您應(yīng)在torch/onnx/symbolic_opset<version>.py中添加符號函數(shù),并按照以下說明進行操作 :

  • torch/onnx/symbolic_opset<version>.py中定義符號功能,例如 torch / onnx / symbolic_opset9.py 。 確保函數(shù)具有與VariableType.h中定義的 ATen 運算符/函數(shù)相同的名稱。
  • 第一個參數(shù)始終是導出的 ONNX 圖。 參數(shù)名稱必須與VariableType.h中的名稱完全匹配,因為分配是通過關(guān)鍵字參數(shù)完成的。
  • 參數(shù)排序不一定與VariableType.h中的匹配,張量(輸入)始終是第一個,然后是非張量參數(shù)。
  • 在符號功能中,如果運算符已經(jīng)在 ONNX 中進行了標準化,我們只需要創(chuàng)建一個節(jié)點即可在圖中表示 ONNX 運算符。
  • 如果輸入?yún)?shù)是張量,但 ONNX 要求標量,則必須顯式進行轉(zhuǎn)換。 輔助函數(shù)_scalar可以將標量張量轉(zhuǎn)換為 python 標量,_if_scalar_type_as可以將 Python 標量轉(zhuǎn)換為 PyTorch 張量。

非 ATen 運營商

如果該運算符是非 ATen 運算符,則必須在相應(yīng)的 PyTorch Function 類中添加符號函數(shù)。 請閱讀以下說明:

  • 在相應(yīng)的 Function 類中創(chuàng)建一個名為symbolic的符號函數(shù)。
  • 第一個參數(shù)始終是導出的 ONNX 圖。
  • 除第一個參數(shù)名稱外,參數(shù)名稱必須與forward中的名稱完全匹配。
  • 輸出元組大小必須與forward的輸出匹配。
  • 在符號功能中,如果運算符已經(jīng)在 ONNX 中進行了標準化,我們只需要創(chuàng)建一個節(jié)點即可在圖中表示 ONNX 運算符。

符號函數(shù)應(yīng)在 Python 中實現(xiàn)。 所有這些功能都與通過 C ++-Python 綁定實現(xiàn)的 Python 方法進行交互,但是直觀地講,它們提供的接口如下所示:

def operator/symbolic(g, *inputs):
  """
  Modifies Graph (e.g., using "op"), adding the ONNX operations representing
  this PyTorch function, and returning a Value or tuple of Values specifying the
  ONNX outputs whose values correspond to the original PyTorch return values
  of the autograd Function (or None if an output is not supported by ONNX).


  Arguments:
    g (Graph): graph to write the ONNX representation into
    inputs (Value...): list of values representing the variables which contain
        the inputs for this function
  """


class Value(object):
  """Represents an intermediate tensor value computed in ONNX."""
  def type(self):
    """Returns the Type of the value."""


class Type(object):
  def sizes(self):
    """Returns a tuple of ints representing the shape of a tensor this describes."""


class Graph(object):
  def op(self, opname, *inputs, **attrs):
    """
    Create an ONNX operator 'opname', taking 'args' as inputs
    and attributes 'kwargs' and add it as a node to the current graph,
    returning the value representing the single output of this
    operator (see the `outputs` keyword argument for multi-return
    nodes).


    The set of operators and the inputs/attributes they take
    is documented at https://github.com/onnx/onnx/blob/master/docs/Operators.md


    Arguments:
        opname (string): The ONNX operator name, e.g., `Abs` or `Add`.
        args (Value...): The inputs to the operator; usually provided
            as arguments to the `symbolic` definition.
        kwargs: The attributes of the ONNX operator, with keys named
            according to the following convention: `alpha_f` indicates
            the `alpha` attribute with type `f`.  The valid type specifiers are
            `f` (float), `i` (int), `s` (string) or `t` (Tensor).  An attribute
            specified with type float accepts either a single float, or a
            list of floats (e.g., you would say `dims_i` for a `dims` attribute
            that takes a list of integers).
        outputs (int, optional):  The number of outputs this operator returns;
            by default an operator is assumed to return a single output.
            If `outputs` is greater than one, this functions returns a tuple
            of output `Value`, representing each output of the ONNX operator
            in positional.
    """

ONNX 圖形 C ++定義在torch/csrc/jit/ir.h中。

這是處理elu運算符缺失的符號函數(shù)的示例。 我們嘗試導出模型,并看到如下錯誤消息:

UserWarning: ONNX export failed on elu because torch.onnx.symbolic_opset9.elu does not exist
RuntimeError: ONNX export failed: Couldn't export operator elu

導出失敗,因為 PyTorch 不支持導出elu運算符。 我們在VariableType.h中找到virtual Tensor elu(const Tensor & input, Scalar alpha, bool inplace) const override;。 這意味著elu是 ATen 運算符。 我們檢查 ONNX 操作員列表,并確認Elu在 ONNX 中已標準化。 我們在symbolic_opset9.py中添加以下行:

def elu(g, input, alpha, inplace=False):
    return g.op("Elu", input, alpha_f=_scalar(alpha))

現(xiàn)在,PyTorch 能夠?qū)С?code>elu運算符。

symbolic_opset9.py 和 symbolic_opset10.py 中還有更多示例。

用于指定操作員定義的界面是實驗性的; 冒險的用戶應(yīng)注意,API 可能會在將來的界面中更改。

定制運算符

按照本教程[使用自定義 C ++運算符擴展 TorchScript ] 之后,您可以在 PyTorch 中創(chuàng)建并注冊自己的自定義 ops 實現(xiàn)。 將這種模型導出到 ONNX 的方法如下:

## Create custom symbolic function
from torch.onnx.symbolic_helper import parse_args
@parse_args('v', 'v', 'f', 'i')
def symbolic_foo_forward(g, input1, input2, attr1, attr2):
    return g.op("Foo", input1, input2, attr1_f=attr1, attr2_i=attr2)


## Register custom symbolic function
from torch.onnx import register_custom_op_symbolic
register_custom_op_symbolic('custom_ops::foo_forward', symbolic_foo_forward, 9)


class FooModel(torch.nn.Module):
    def __init__(self, attr1, attr2):
        super(FooModule, self).__init__()
        self.attr1 = attr1
        self.attr2 = attr2


    def forward(self, input1, input2):
        # Calling custom op
        return torch.ops.custom_ops.foo_forward(input1, input2, self.attr1, self.attr2)


model = FooModel(attr1, attr2)
torch.onnx.export(model, (dummy_input1, dummy_input2), 'model.onnx')

根據(jù)自定義運算符的不同,您可以將其導出為現(xiàn)有 ONNX 操作之一或組合。 您也可以將其導出為 ONNX 中的自定義操作。 在這種情況下,您將需要通過匹配的自定義操作實現(xiàn)來擴展選擇的后端,例如 Caffe2 定制操作, ONNX Runtime 定制操作。

常見問題解答

問:我已經(jīng)導出了我的 lstm 模型,但是它的輸入大小似乎是固定的?

跟蹤器將示例輸入形狀記錄在圖中。 如果模型應(yīng)接受動態(tài)形狀的輸入,則可以在導出 api 中使用參數(shù) <cite&dynamic_axes</cite& 。


layer_count = 4

model = nn.LSTM(10, 20, num_layers=layer_count, bidirectional=True)
model.eval()

with torch.no_grad():
input = torch.randn(5, 3, 10)
h0 = torch.randn(layer_count * 2, 3, 20)
c0 = torch.randn(layer_count * 2, 3, 20)
output, (hn, cn) = model(input, (h0, c0))

# default export
torch.onnx.export(model, (input, (h0, c0)), 'lstm.onnx')
onnx_model = onnx.load('lstm.onnx')
# input shape [5, 3, 10]
print(onnx_model.graph.input[0])

# export with `dynamic_axes`
torch.onnx.export(model, (input, (h0, c0)), 'lstm.onnx',
input_names=['input', 'h0', 'c0'],
output_names=['output', 'hn', 'cn'],
dynamic_axes={'input': {0: 'sequence'}, 'output': {0: 'sequence'}})
onnx_model = onnx.load('lstm.onnx')
# input shape ['sequence', 3, 10]
print(onnx_model.graph.input[0])

問:如何導出帶有循環(huán)的模型?

請簽出跟蹤與腳本編寫。

問:ONNX 是否支持隱式標量數(shù)據(jù)類型轉(zhuǎn)換?

不,但是出口商將嘗試處理該部分。 標量在 ONNX 中轉(zhuǎn)換為恒定張量。 導出器將嘗試找出標量的正確數(shù)據(jù)類型。 但是,對于無法執(zhí)行此操作的情況,您將需要手動提供數(shù)據(jù)類型信息。 這通常發(fā)生在腳本模型中,其中未記錄數(shù)據(jù)類型。 我們正在嘗試改進數(shù)據(jù)類型在導出器中的傳播,以便將來不再需要手動更改。


class ImplicitCastType(torch.jit.ScriptModule):
@torch.jit.script_method
def forward(self, x):
# Exporter knows x is float32, will export '2' as float32 as well.
y = x + 2
# Without type propagation, exporter doesn't know the datatype of y.
# Thus '3' is exported as int64 by default.
return y + 3
# The following will export correctly.
# return y + torch.tensor([3], dtype=torch.float32)

x = torch.tensor([1.0], dtype=torch.float32)
torch.onnx.export(ImplicitCastType(), x, 'models/implicit_cast.onnx',
example_outputs=ImplicitCastType()(x))

功能

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 格式。 這個導出器運行一次您的模型,以便跟蹤要導出的模型執(zhí)行情況。 目前,它支持一組有限的動態(tài)模型(例如 RNN)。

參數(shù)

  • 模型 (torch.nn.Module)–要導出的模型。

  • 參數(shù)(參數(shù)元組)–模型的輸入,例如,使得model(*args)是模型的有效調(diào)用。 任何非 Tensor 參數(shù)將被硬編碼到導出的模型中; 任何 Tensor 參數(shù)將按照在 args 中出現(xiàn)的順序成為導出模型的輸入。 如果 args 是一個 Tensor,則相當于用該 Tensor 的 1 元元組調(diào)用了它。 (注意:當前不支持將關(guān)鍵字參數(shù)傳遞給模型。如果需要,請給我們喊叫。)

  • f –類似于文件的對象(必須實現(xiàn)返回文件描述符的 fileno)或包含文件名的字符串。 二進制 Protobuf 將被寫入此文件。

  • export_params (布爾 默認為 True )–如果指定,將導出所有參數(shù)。 如果要導出未經(jīng)訓練的模型,請將其設(shè)置為 False。 在這種情況下,導出的模型將首先以其所有參數(shù)作為參數(shù),順序由model.state_dict().values()指定

  • 詳細 (bool 默認為 False )–如果指定,我們將打印出導出跟蹤的調(diào)試描述。

  • 訓練 (bool , 默認為 False )–以訓練模式導出模型。 目前,ONNX 僅面向?qū)С瞿P鸵赃M行推理,因此通常不需要將其設(shè)置為 True。

  • input_names (字符串列表 , 默認空列表)–依次分配給圖形輸入節(jié)點的名稱

  • output_names (字符串列表 , 默認空列表)–依次分配給圖形輸出節(jié)點的名稱

  • (bool 默認為 False )– [不推薦使用。 使用 operator_export_type]以 aten 模式導出模型。 如果使用 aten 模式,則 symbolic_opset <版本> .py 中的函數(shù)所導出的所有 ops 原始文件都將作為 ATen ops 導出。

  • export_raw_ir (布爾 , 默認為 False )– [不建議使用。 使用 operator_export_type]直接導出內(nèi)部 IR,而不是將其轉(zhuǎn)換為 ONNX ops。

  • operator_export_type (枚舉 , 默認 OperatorExportTypes.ONNX )– OperatorExportTypes.ONNX:所有操作均作為常規(guī) ONNX 操作導出。 OperatorExportTypes.ONNX_ATEN:所有操作均導出為 ATen 操作。 OperatorExportTypes.ONNX_ATEN_FALLBACK:如果缺少符號,請使用 ATen op。 OperatorExportTypes.RAW:導出原始 ir。

  • opset_version (python:int , 默認為 9 )–默認情況下,我們將模型導出到 onnx 子模塊的 opset 版本。 由于 ONNX 的最新 opset 可能會在下一個穩(wěn)定版本之前發(fā)展,因此默認情況下,我們會導出到一個穩(wěn)定的 opset 版本。 目前,受支持的穩(wěn)定 opset 版本為 9。opset_version 必須為 _onnx_master_opset 或在 torch / onnx / symbolic_helper.py 中定義的 _onnx_stable_opsets 中。

  • do_constant_folding (bool , 默認 False )–如果為 True,則在導出期間將恒定折疊優(yōu)化應(yīng)用于模型。 常量折疊優(yōu)化將用預(yù)先計算的常量節(jié)點替換一些具有所有常量輸入的操作。

  • example_outputs (張量元組 , 默認無)–導出 ScriptModule 或 TorchScript 函數(shù)時必須提供 example_outputs。

  • strip_doc_string (bool , 默認 True )–如果為 True,則從導出的模型中刪除字段“ doc_string”,有關(guān) 堆棧跟蹤。

  • example_outputs –正在導出的模型的示例輸出。

  • dynamic_axes (dict <字符串 , dict < python:int , 字符串 > > dict <字符串 , 列表 ( python:int > , 默認為空字典)–

一個字典,用于指定輸入/輸出的動態(tài)軸,例如:-KEY:輸入和/或輸出名稱-VALUE:給定鍵的動態(tài)軸的索引,以及可能用于導出動態(tài)軸的名稱。 通常,該值是根據(jù)以下方式之一或兩者的組合定義的:(1)。 指定提供的輸入的動態(tài)軸的整數(shù)列表。 在這種情況下,將在導出過程中自動生成名稱并將其應(yīng)用于提供的輸入/輸出的動態(tài)軸。 或(2)。 一個內(nèi)部字典,該字典指定從對應(yīng)的輸入/輸出中的動態(tài)軸的索引到在導出過程中希望在此輸入/輸出的該軸上應(yīng)用的名稱的映射。

例。 如果我們的輸入和輸出具有以下形狀:

  shape(input_1) = ('b', 3, 'w', 'h')
  and shape(input_2) = ('b', 4)
  and shape(output)  = ('b', 'd', 5)

  Then dynamic axes can be defined either as:

  

  (a). ONLY INDICES:

  

dynamic_axes = {'input_1':[0,2,3],'input_2':[0],'output':[0,1]}

其中將為導出的動態(tài)軸生成自動名稱

  (b). INDICES WITH CORRESPONDING NAMES:

  

dynamic_axes = {'input_1':{0:'batch',1:'width',2:'height'},'input_2':{0:'batch'},'output':{0:'batch', 1:“檢測”}

提供的名稱將應(yīng)用于導出的動態(tài)軸

  (c). MIXED MODE OF (a) and (b)

  

dynamic_axes = {'input_1':[0,2,3],'input_2':{0:'batch'},'output':[0,1]}

  • keep_initializers_as_inputs (bool , 默認值 None )–如果為 True,則導出的圖中的所有初始化程序(通常對應(yīng)于參數(shù))也將 被添加為圖形的輸入。 如果為 False,則不會將初始化程序添加為圖形的輸入,而僅將非參數(shù)輸入添加為輸入。 通過執(zhí)行這些圖形的后端/運行時,這可以允許進行更好的優(yōu)化(例如恒定折疊等)。 如果未指定(默認為“無”),則按以下方式自動選擇行為。 如果 operator_export_type 為 OperatorExportTypes.ONNX,則該行為等效于將此參數(shù)設(shè)置為 False。 對于 operator_export_type 的其他值,此行為等同于將此參數(shù)設(shè)置為 True。 請注意,對于 ONNX opset 版本< 9,初始化器必須是圖形輸入的一部分。 因此,如果 opset_version 參數(shù)設(shè)置為 8 或更低,則該參數(shù)將被忽略。

torch.onnx.register_custom_op_symbolic(symbolic_name, symbolic_fn, opset_version)?

torch.onnx.operators.shape_as_tensor(x)?

torch.onnx.set_training(model, mode)?

上下文管理器將“模型”的訓練模式臨時設(shè)置為“模式”,當我們退出 with 塊時將其重置。 如果模式為“無”,則為無操作。

torch.onnx.is_in_onnx_export()?

檢查它是否在 ONNX 導出的中間。 此函數(shù)在 torch.onnx.export()的中間返回 True。 torch.onnx.export 應(yīng)該使用單線程執(zhí)行。

以上內(nèi)容是否對您有幫助:
在線筆記
App下載
App下載

掃描二維碼

下載編程獅App

公眾號
微信公眾號

編程獅公眾號