【推理引擎】ONNX 模型解析

虔誠的樹發表於2022-03-27
模型

定義模型結構

首先使用 PyTorch 定義一個簡單的網路模型：

class ConvBnReluBlock(nn.Module):
    def __init__(self) -> None:
        super().__init__()

        self.conv1 = nn.Conv2d(3, 64, 3)
        self.bn1 = nn.BatchNorm2d(64)
        self.maxpool1 = nn.MaxPool2d(3, 1)

        self.conv2 = nn.Conv2d(64, 32, 3)
        self.bn2 = nn.BatchNorm2d(32)
    
        self.relu = nn.ReLU()

    def forward(self, x):
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.maxpool1(out)
        
        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)
        
        return out

在匯出模型之前，需要提前定義一些變數：

model = ConvBnReluBlock()     # 定義模型物件
x = torch.randn(2, 3, 255, 255)      # 定義輸入張量

然後使用 PyTorch 官方 API（torch.onnx.export）匯出 ONNX 格式的模型：

# way1:
torch.onnx.export(model, (x), "conv_bn_relu_evalmode.onnx", input_names=["input"], output_names=['output'])

# way2:
import torch._C as _C
TrainingMode = _C._onnx.TrainingMode
torch.onnx.export(model, (x), "conv_bn_relu_trainmode.onnx", input_names=["input"], output_names=['output'],
                opset_version=12,                    # 預設版本為9,但是如果低於12，將不能正確匯出 Dropout 和 BatchNorm 節點
                training=TrainingMode.TRAINING,      # 預設模式為 TrainingMode.EVAL
                do_constant_folding=False)           # 常量摺疊，預設為True，但是如果使用TrainingMode.TRAINING模式，則需要將其關閉

# way3
torch.onnx.export(model,
                (x),
                "conv_bn_relu_dynamic.onnx",
                input_names=['input'],
                output_names=['output'],
                dynamic_axes={'input': {0: 'batch_size', 2: 'input_width', 3: 'input_height'},
                            'output': {0: 'batch_size', 2: 'output_width', 3: 'output_height'}})

可以看到，這裡主要以三種方式匯出模型，下面分別介紹區別：

  • way1：如果模型中存在 BatchNorm 或者 Dropout，我們在匯出模型前會首先將其設定成 eval 模式，但是這裡我們即使忘記設定也無所謂，因為在匯出模型時會自動設定（export函式中training引數的預設值為TrainingMode.EVAL）。
  • way2：如果我們想匯出完整的模型結構，包括 BatchNorm 和 Dropout，則應該將 training 屬性設定為 train 模式。
  • way3：如果想要匯出動態輸入的模型結構，則需要設定 dynamic_axes 屬性，比如這裡我們將第一、三和四維設定成動態結構，那麼我們就可以輸入任何Batch大小、任何長寬尺度的RGB影像。

下圖分別將這三種匯出方式的模型結構使用 Netron 視覺化：

分析模型結構

這裡參考了BBuf大佬的講解：【傳送門：https://zhuanlan.zhihu.com/p/346511883】
接下來主要針對 way1 方式匯出的ONNX模型進行深入分析。

ONNX格式定義：https://github.com/onnx/onnx/blob/master/onnx/onnx.proto
在這個檔案中，定義了多個核心物件：ModelProto、GraphProto、NodeProto、ValueInfoProto、TensorProto 和 AttributeProto。

在載入ONNX模型之後，就獲得了一個ModelProto，其中包含一些

  • 版本資訊(本例中：ir_version = 7)
  • 生成者資訊：producer_name: pytorch，producer_version: 1.10，這兩個屬性主要用來說明由哪些框架哪個版本匯出的onnx
  • 核心元件：GraphProto

在 GraphProto 中，有如下幾個屬性需要注意：

  • name：本例中：name = 'torch-jit-export'
  • input 陣列：
    [name: "input"
type {
  tensor_type {
    elem_type: 1
    shape {
      dim {
        dim_value: 2
      }
      dim {
        dim_value: 3
      }
      dim {
        dim_value: 255
      }
      dim {
        dim_value: 255
      }
    }
  }
}
]
  • output 陣列：
    [name: "output"
type {
  tensor_type {
    elem_type: 1
    shape {
      dim {
        dim_value: 2
      }
      dim {
        dim_value: 32
      }
      dim {
        dim_value: 249
      }
      dim {
        dim_value: 249
      }
    }
  }
}
]
  • node 陣列，該陣列中包含了模型中所有的計算節點（本例中："Conv_0"、"Relu_1"、"MaxPool_2"、"Conv_3"、"Relu_4"），以及各個節點的屬性，：
     [input: "input"
input: "23"
input: "24"
output: "22"
name: "Conv_0"
op_type: "Conv"
attribute {
  name: "dilations"
  ints: 1
  ints: 1
  type: INTS
}
attribute {
  name: "group"
  i: 1
  type: INT
}
attribute {
  name: "kernel_shape"
  ints: 3
  ints: 3
  type: INTS
}
attribute {
  name: "pads"
  ints: 0
  ints: 0
  ints: 0
  ints: 0
  type: INTS
}
attribute {
  name: "strides"
  ints: 1
  ints: 1
  type: INTS
}
, 
input: "22"
output: "17"
name: "Relu_1"
op_type: "Relu"
, input: "17"
output: "18"
name: "MaxPool_2"
op_type: "MaxPool"
attribute {
  name: "kernel_shape"
  ints: 3
  ints: 3
  type: INTS
}
attribute {
  name: "pads"
  ints: 0
  ints: 0
  ints: 0
  ints: 0
  type: INTS
}
attribute {
  name: "strides"
  ints: 1
  ints: 1
  type: INTS
}
, 
input: "18"
input: "26"
input: "27"
output: "25"
name: "Conv_3"
op_type: "Conv"
attribute {
  name: "dilations"
  ints: 1
  ints: 1
  type: INTS
}
attribute {
  name: "group"
  i: 1
  type: INT
}
attribute {
  name: "kernel_shape"
  ints: 3
  ints: 3
  type: INTS
}
attribute {
  name: "pads"
  ints: 0
  ints: 0
  ints: 0
  ints: 0
  type: INTS
}
attribute {
  name: "strides"
  ints: 1
  ints: 1
  type: INTS
}
, 
input: "25"
output: "output"
name: "Relu_4"
op_type: "Relu"
]
    通過以上 node 的輸入輸出資訊，可提取出節點之間的拓撲關係，構建出一個完整的神經網路。
  • initializer 陣列：存放模型的權重引數。
    [dims: 64
dims: 3
dims: 3
dims: 3
data_type: 1
name: "23"
raw_data: "\220\251\001>\232\326&>\253\227\372 ... 省略一眼望不到頭的內容 ... "

dims: 64
data_type: 1
name: "24"
raw_data: "Rt\347\275\005\203\0 ..."

dims: 32
dims: 64
dims: 3
dims: 3
data_type: 1
name: "26"
raw_data: "9\022\273;+^\004\2 ..."

...

至此，我們已經分析完 GraphProto 的內容，下面根據圖中的一個節點視覺化說明以上內容：

從圖中可以發現，Conv 節點的輸入包含三個部分：輸入的影像（input）、權重（這裡以數字23代表該節點權重W的名字）以及偏置（這裡以數字24表示該節點偏置B的名字）；輸出內容的名字為22；屬性資訊包括dilations、group、kernel_shape、pads和strides，不同節點會具有不同的屬性資訊。在initializer陣列中，我們可以找到該Conv節點權重（name:23）對應的值(raw_data)，並且可以清楚地看到維度資訊（64X3X3X3）。

相關文章