PyTorch torch.onnx

2020-09-15 10:47 更新

原文: PyTorch torch.onnx

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

這是一個(gè)簡單的腳本,可以將 Torchvision 中定義的經(jīng)過預(yù)訓(xùn)練的 AlexNet 導(dǎo)出到 ONNX 中。 它運(yùn)行一輪推斷,然后將生成的跟蹤模型保存到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是二進(jìn)制 protobuf 文件,其中包含您導(dǎo)出的模型的網(wǎng)絡(luò)結(jié)構(gòu)和參數(shù)(在本例中為 AlexNet)。 關(guān)鍵字參數(shù)verbose=True使導(dǎo)出程序打印出人類可讀的網(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 庫來驗(yàn)證 protobuf。 您可以使用 conda 安裝ONNX

conda install -c conda-forge onnx

然后,您可以運(yùn)行:

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 運(yùn)行導(dǎo)出的腳本,您將需要安裝 <cite>caffe2</cite> :如果尚未安裝,請按照安裝說明進(jìn)行操作。

一旦安裝了這些,就可以將后端用于 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 運(yùn)行導(dǎo)出的模型,您將需要安裝 <cite>ONNX Runtime</cite> :請按照這些說明進(jìn)行操作。

一旦安裝了這些,就可以將后端用于 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 模型導(dǎo)出到 ONNX 的另一本教程。 。

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

跟蹤與腳本編寫

ONNX 導(dǎo)出器可以是基于跟蹤的和基于腳本的導(dǎo)出器。

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

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

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)

使用基于跟蹤的導(dǎo)出器,我們得到結(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)

為了利用基于腳本的導(dǎo)出器捕獲動態(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)在,導(dǎo)出的 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)范圍的后端進(jìn)行驗(yà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)]

局限性

  • 導(dǎo)出中目前不支持張量就地索引分配,例如 <cite>data [index] = new_data</cite> 。 解決此類問題的一種方法是使用運(yùn)算符<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 中沒有張量列表的概念。 沒有這個(gè)概念,很難導(dǎo)出消耗或產(chǎn)生張量列表的運(yùn)算符,尤其是在導(dǎo)出時(shí)不知道張量列表的長度的情況下。

  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 輸入/輸出支持。 也接受字典和字符串,但不建議使用它們。 用戶需要仔細(xì)驗(yàn)證自己的字典輸入,并記住動態(tài)查詢不可用。

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

支持的運(yùn)營商

支持以下運(yùn)算符:

  • 批量標(biāo)準(zhǔn)
  • ConstantPadNd
  • 轉(zhuǎn)換
  • 退出
  • 嵌入(不支持可選參數(shù))
  • FeatureDropout(不支持訓(xùn)練模式)
  • 指數(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
  • 細(xì)胞
  • 最大鉗位
  • 最小鉗位
  • 康卡特
  • cos
  • cumsum
  • 暗淡的
  • div
  • 退出
  • lu
  • 空的
  • 空的喜歡
  • 當(dāng)量
  • 埃爾夫
  • 經(jīng)驗(yàn)值
  • 擴(kuò)大
  • 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ù)表達(dá)式
  • lt
  • masked_fill
  • 最大值
  • 意思
  • 毫米
  • 多項(xiàng)式
  • 狹窄
  • NE
  • 負(fù)數(shù)
  • 非零
  • 規(guī)范
  • 那些
  • 喜歡
  • 要么
  • 置換
  • pixel_shuffle
  • 戰(zhàn)俘
  • prelu(不支持輸入通道之間共享的單個(gè)權(quán)重)
  • 產(chǎn)品
  • 蘭德
  • 蘭德
  • randn_like
  • 倒數(shù)
  • Reflection_pad
  • 露露
  • 重復(fù)
  • 復(fù)制墊
  • 重塑
  • reshape_as
  • 回合
  • 雷雷魯
  • rsqrt
  • 訂閱
  • 分散
  • scatter_add
  • 選擇
  • 塞盧
  • 乙狀結(jié)腸
  • 標(biāo)志
  • 尺寸
  • 切片
  • 軟最大
  • 軟加
  • 分類
  • 分裂
  • sqrt
  • 性病
  • 子(不支持非零 Alpha)
  • ?
  • 棕褐色
  • 閾值(不支持非零閾值/非零值)
  • 托普
  • 轉(zhuǎn)置
  • type_as
  • 展開(與 ATen-Caffe2 集成的實(shí)驗(yàn)支持)
  • 獨(dú)特
  • 松開
  • upsample_nearest1d
  • upsample_nearest2d
  • upsample_nearest3d
  • 視圖
  • 哪里
  • zeros_like

上面設(shè)置的運(yùn)算符足以導(dǎo)出以下模型:

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

添加對運(yùn)營商的支持

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

ATen 運(yùn)算符

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

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

非 ATen 運(yùn)營商

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

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

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

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運(yùn)算符缺失的符號函數(shù)的示例。 我們嘗試導(dǎo)出模型,并看到如下錯(cuò)誤消息:

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

導(dǎo)出失敗,因?yàn)?PyTorch 不支持導(dǎo)出elu運(yùn)算符。 我們在VariableType.h中找到virtual Tensor elu(const Tensor & input, Scalar alpha, bool inplace) const override;。 這意味著elu是 ATen 運(yùn)算符。 我們檢查 ONNX 操作員列表,并確認(rèn)Elu在 ONNX 中已標(biāo)準(zhǔn)化。 我們在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運(yùn)算符。

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

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

定制運(yùn)算符

按照本教程[使用自定義 C ++運(yùn)算符擴(kuò)展 TorchScript ] 之后,您可以在 PyTorch 中創(chuàng)建并注冊自己的自定義 ops 實(shí)現(xiàn)。 將這種模型導(dǎo)出到 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ù)自定義運(yùn)算符的不同,您可以將其導(dǎo)出為現(xiàn)有 ONNX 操作之一或組合。 您也可以將其導(dǎo)出為 ONNX 中的自定義操作。 在這種情況下,您將需要通過匹配的自定義操作實(shí)現(xiàn)來擴(kuò)展選擇的后端,例如 Caffe2 定制操作, ONNX Runtime 定制操作。

常見問題解答

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

跟蹤器將示例輸入形狀記錄在圖中。 如果模型應(yīng)接受動態(tài)形狀的輸入,則可以在導(dǎo)出 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])

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

請簽出跟蹤與腳本編寫。

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

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


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)?

將模型導(dǎo)出為 ONNX 格式。 這個(gè)導(dǎo)出器運(yùn)行一次您的模型,以便跟蹤要導(dǎo)出的模型執(zhí)行情況。 目前,它支持一組有限的動態(tài)模型(例如 RNN)。

參數(shù)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  • example_outputs –正在導(dǎo)出的模型的示例輸出。

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

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

其中將為導(dǎo)出的動態(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)用于導(dǎo)出的動態(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 , 默認(rèn)值 None )–如果為 True,則導(dǎo)出的圖中的所有初始化程序(通常對應(yīng)于參數(shù))也將 被添加為圖形的輸入。 如果為 False,則不會將初始化程序添加為圖形的輸入,而僅將非參數(shù)輸入添加為輸入。 通過執(zhí)行這些圖形的后端/運(yùn)行時(shí),這可以允許進(jìn)行更好的優(yōu)化(例如恒定折疊等)。 如果未指定(默認(rèn)為“無”),則按以下方式自動選擇行為。 如果 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)?

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

torch.onnx.is_in_onnx_export()?

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

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

掃描二維碼

下載編程獅App

公眾號
微信公眾號

編程獅公眾號