TensorRT ONNX 基礎

tensorRT從零起步邁向高性能工業級部署（就業導向）課程筆記，講師講的不錯，可以去看原視頻支持下。

概述

TensorRT 的核心在于對模型算子的優化（合并算子、利用當前 GPU 特性選擇特定的核函數等多種策略），通過 TensorRT，能夠在 Nvidia 系列 GPU 上獲得最好的性能。
TensorRT 模型需要在目標 GPU 上以實際運行的方式選擇最優的算法和配置（不同的 GPU 的許多特性的不一樣，在特定 GPU 上跑一跑，再知道怎樣最快）。
也因此 TensorRT 得到的模型只能在特定的環境下運行（編譯時的 TensorRT 版本、CUDA 版本、GPU 型號等）。如果不是在完全相同的環境下運行 TensorRT 引擎來推理有時是直接無法運行的、有時是可以運行，但是并不能保證是最佳的，因此盡量不要這么做。
本節主要知識點
- TensorRT 模型定義的方式
- 編譯過程配置
- 推理過程實現
- 插件實現
- onnx 理解

下圖展示了 TensoRT 優化前后的模型，TensorRT 會找到一些可以合并、優化的算子，進行合并。

在這里插入圖片描述

TensorRT 模型描述方案選擇

最底層：TensorRT C++接口、Python 接口
常見的工作流：uff, onnx, caffe

在這里插入圖片描述

tensorrtx，這是一個 github 上第三方的庫，在官方接口的基礎上封裝了 ResNet 等常用的網絡模型。
本課程的選擇：PyTorch -> ONNX -> TensorRT

選擇 ONNX 的一個好處是：ONNX 是一個通用的網絡模型的中間格式，熟悉了 ONNX 格式之后，不僅是轉到 TensorRT 引擎，如果后續有其他需要也可以方便地轉換到其他推理引擎如 ncnn、mnn 等
TensorRT 的一般需要包含的頭文件是 NvInfer.h 和 NvInferRuntime.h，而 TensorRT 的庫文件一覽如下：

在這里插入圖片描述

TensorRT C++ 基本接口模型構建

下面的代碼通過一個最簡單的網絡展示了 TensorRT C++ 一些基本接口來構建一個模型的過程。

// tensorRT include
#include <NvInfer.h>
#include <NvInferRuntime.h>// cuda include
#include <cuda_runtime.h>// system include
#include <stdio.h>class TRTLogger : public nvinfer1::ILogger{
public:virtual void log(Severity severity, nvinfer1::AsciiChar const* msg) noexcept override{if(severity <= Severity::kVERBOSE){printf("%d: %s\n", severity, msg);}}
};nvinfer1::Weights make_weights(float* ptr, int n){nvinfer1::Weights w;w.count = n;w.type = nvinfer1::DataType::kFLOAT;w.values = ptr;return w;
}int main(){// 本代碼主要實現一個最簡單的神經網絡 figure/simple_fully_connected_net.png TRTLogger logger; // logger是必要的，用來捕捉warning和info等// ----------------------------- 1. 定義 builder, config 和network -----------------------------// 這是基本需要的組件//形象的理解是你需要一個builder去build這個網絡，網絡自身有結構，這個結構可以有不同的配置nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger);// 創建一個構建配置，指定TensorRT應該如何優化模型，tensorRT生成的模型只能在特定配置下運行nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();// 創建網絡定義，其中createNetworkV2(1)表示采用顯性batch size，新版tensorRT(>=7.0)時，不建議采用0非顯性batch size// 因此貫穿以后，請都采用createNetworkV2(1)而非createNetworkV2(0)或者createNetworknvinfer1::INetworkDefinition* network = builder->createNetworkV2(1);// 構建一個模型/*Network definition:image|linear (fully connected)  input = 3, output = 2, bias = True     w=[[1.0, 2.0, 0.5], [0.1, 0.2, 0.5]], b=[0.3, 0.8]|sigmoid|prob*/// ----------------------------- 2. 輸入，模型結構和輸出的基本信息 -----------------------------const int num_input = 3;   // in_channelconst int num_output = 2;  // out_channelfloat layer1_weight_values[] = {1.0, 2.0, 0.5, 0.1, 0.2, 0.5}; // 前3個給w1的rgb，后3個給w2的rgb float layer1_bias_values[]   = {0.3, 0.8};//輸入指定數據的名稱、數據類型和完整維度，將輸入層添加到網絡nvinfer1::ITensor* input = network->addInput("image", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4(1, num_input, 1, 1));nvinfer1::Weights layer1_weight = make_weights(layer1_weight_values, 6);nvinfer1::Weights layer1_bias   = make_weights(layer1_bias_values, 2);//添加全連接層auto layer1 = network->addFullyConnected(*input, num_output, layer1_weight, layer1_bias);      // 注意對input進行了解引用//添加激活層 auto prob = network->addActivation(*layer1->getOutput(0), nvinfer1::ActivationType::kSIGMOID); // 注意更嚴謹的寫法是*(layer1->getOutput(0)) 即對getOutput返回的指針進行解引用// 將我們需要的prob標記為輸出network->markOutput(*prob->getOutput(0));printf("Workspace Size = %.2f MB\n", (1 << 28) / 1024.0f / 1024.0f); // 256Mibconfig->setMaxWorkspaceSize(1 << 28);builder->setMaxBatchSize(1); // 推理時 batchSize = 1 // ----------------------------- 3. 生成engine模型文件 -----------------------------//TensorRT 7.1.0版本已棄用buildCudaEngine方法，統一使用buildEngineWithConfig方法nvinfer1::ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);if(engine == nullptr){printf("Build engine failed.\n");return -1;}// ----------------------------- 4. 序列化模型文件并存儲 -----------------------------// 將模型序列化，并儲存為文件nvinfer1::IHostMemory* model_data = engine->serialize();FILE* f = fopen("engine.trtmodel", "wb");fwrite(model_data->data(), 1, model_data->size(), f);fclose(f);// 卸載順序按照構建順序倒序model_data->destroy();engine->destroy();network->destroy();config->destroy();builder->destroy();printf("Done.\n");return 0;
}

重點提煉：

必須使用 createNetworkV2 ，并制定 1（表示顯性 batch），createNetwork 已經廢棄，非顯性 batch 官方不推薦。這個直接影響到推理時是 enqueue 還是 enqueueV2
builder、config 等指針，記得釋放，使用 ptr->destroy()，否則會有內存泄漏
markOutput 表示是該模型的輸出節點，mark 幾次，就有幾個輸出，addInput 幾次，就有幾個輸入，這與推理時的輸入輸出相呼應
workSpaceSize 是工作空間的大小，某些 layer 需要使用額外存儲時，不會自己分配空間，而是為了內存復用，直接找 TensorRT 要 workspace 空間，指的是這個意思
一定要記住，保存的模型只能適應編譯時的 TensorRT版本、CUDA 版本、GPU 型號等環境。也只能保證在完全相同的配置時是最優的。如果模型跨不同設備執行，有時也可以運行，但是不是最優的，也不推薦。

TensorRT C++ 基本接口模型推理

下面的代碼對上一小節構建的簡單網絡進行推理。


// tensorRT include
#include <NvInfer.h>
#include <NvInferRuntime.h>// cuda include
#include <cuda_runtime.h>// system include
#include <stdio.h>
#include <math.h>#include <iostream>
#include <fstream>
#include <vector>using namespace std;
// 上一節的代碼class TRTLogger : public nvinfer1::ILogger{
public:virtual void log(Severity severity, nvinfer1::AsciiChar const* msg) noexcept override{if(severity <= Severity::kINFO){printf("%d: %s\n", severity, msg);}}
} logger;nvinfer1::Weights make_weights(float* ptr, int n){nvinfer1::Weights w;w.count = n;w.type = nvinfer1::DataType::kFLOAT;w.values = ptr;return w;
}bool build_model(){// 這里的build_model函數即是做了和上面構建模型小節一樣的事情，不再贅述
}vector<unsigned char> load_file(const string& file){ifstream in(file, ios::in | ios::binary);if (!in.is_open())return {};in.seekg(0, ios::end);size_t length = in.tellg();std::vector<uint8_t> data;if (length > 0){in.seekg(0, ios::beg);data.resize(length);in.read((char*)&data[0], length);}in.close();return data;
}void inference(){// ------------------------------ 1. 準備模型并加載   ----------------------------TRTLogger logger;auto engine_data = load_file("engine.trtmodel");// 執行推理前，需要創建一個推理的runtime接口實例。與builer一樣，runtime需要logger：nvinfer1::IRuntime* runtime   = nvinfer1::createInferRuntime(logger);// 將模型從讀取到engine_data中，則可以對其進行反序列化以獲得enginenvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(engine_data.data(), engine_data.size());if(engine == nullptr){printf("Deserialize cuda engine failed.\n");runtime->destroy();return;}nvinfer1::IExecutionContext* execution_context = engine->createExecutionContext();cudaStream_t stream = nullptr;// 創建CUDA流，以確定這個batch的推理是獨立的cudaStreamCreate(&stream);/*Network definition:image|linear (fully connected)  input = 3, output = 2, bias = True     w=[[1.0, 2.0, 0.5], [0.1, 0.2, 0.5]], b=[0.3, 0.8]|sigmoid|prob*/// ------------------------------ 2. 準備好要推理的數據并搬運到GPU   ----------------------------float input_data_host[] = {1, 2, 3};float* input_data_device = nullptr;float output_data_host[2];float* output_data_device = nullptr;cudaMalloc(&input_data_device, sizeof(input_data_host));cudaMalloc(&output_data_device, sizeof(output_data_host));cudaMemcpyAsync(input_data_device, input_data_host, sizeof(input_data_host), cudaMemcpyHostToDevice, stream);// 用一個指針數組指定input和output在gpu中的指針。float* bindings[] = {input_data_device, output_data_device};// ------------------------------ 3. 推理并將結果搬運回CPU   ----------------------------bool success      = execution_context->enqueueV2((void**)bindings, stream, nullptr);cudaMemcpyAsync(output_data_host, output_data_device, sizeof(output_data_host), cudaMemcpyDeviceToHost, stream);cudaStreamSynchronize(stream);printf("output_data_host = %f, %f\n", output_data_host[0], output_data_host[1]);// ------------------------------ 4. 釋放內存 ----------------------------printf("Clean memory\n");cudaStreamDestroy(stream);execution_context->destroy();engine->destroy();runtime->destroy();// ------------------------------ 5. 手動推理進行驗證 ----------------------------const int num_input = 3;const int num_output = 2;float layer1_weight_values[] = {1.0, 2.0, 0.5, 0.1, 0.2, 0.5};float layer1_bias_values[]   = {0.3, 0.8};printf("手動驗證計算結果：\n");for(int io = 0; io < num_output; ++io){float output_host = layer1_bias_values[io];for(int ii = 0; ii < num_input; ++ii){output_host += layer1_weight_values[io * num_input + ii] * input_data_host[ii];}// sigmoidfloat prob = 1 / (1 + exp(-output_host));printf("output_prob[%d] = %f\n", io, prob);}
}int main(){if(!build_model()){return -1;}inference();return 0;
}

重點提煉：

bindings 是對 TensorRT 輸入輸出張量的描述，bindings = input_tensor + output_tensor，比如 input 有 a，output 有 b, c, d，那么 bindings = [a, b, c, d]，可以就當成個數組：bindings[0] = a，bindings[2] = c。獲取 bindings：engine->getBindingDimensions(0)
enqueueV2 是異步推理，加入到 stream 隊列等待執行，輸入的 bindings 則是 tensor 指針（注意是 device pointer）。其 shape 對應于編譯時指定的輸入輸出的 shape（目前演示的shape都是靜態的）
createExecutionContext 可以執行多次，允許一個引擎具有多個執行上下文，不過看看就好，別當真。

動態shape

動態 shape，即在構建模型時可以先不確定 shape，而是指定一個動態范圍： $[L ? H]$ ，推理時再確定 shape，允許的范圍即是： $L<=shape<=HL\ <=\ shape\ <=\ H$

在構建時，主要是在這幾行代碼，指定 shape 的動態范圍，其他與之前類似：

int maxBatchSize = 10;
printf("Workspace Size = %.2f MB\n", (1 << 28) / 1024.0f / 1024.0f);
// 配置暫存存儲器，用于layer實現的臨時存儲，也用于保存中間激活值
config->setMaxWorkspaceSize(1 << 28);// --------------------------------- 2.1 關于profile ----------------------------------
// 如果模型有多個輸入，則必須多個profile
auto profile = builder->createOptimizationProfile();// 配置最小允許1 x 1 x 3 x 3
profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMIN, nvinfer1::Dims4(1, num_input, 3, 3));
profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kOPT, nvinfer1::Dims4(1, num_input, 3, 3));// 配置最大允許10 x 1 x 5 x 5
// if networkDims.d[i] != -1, then minDims.d[i] == optDims.d[i] == maxDims.d[i] == networkDims.d[i]
profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMAX, nvinfer1::Dims4(maxBatchSize, num_input, 5, 5));
config->addOptimizationProfile(profile);

在推理時，增加的是這些，來指定具體的 shape：

int ib = 2;
int iw = 3;
int ih = 3;// 明確當前推理時，使用的數據輸入大小
execution_context->setBindingDimensions(0, nvinfer1::Dims4(ib, 1, ih, iw));

重點提煉：

OptimizationProfile 是一個優化配置文件，用來指定輸入的 shape 可以變化的動態范圍
如果 ONNX 的某個維度是 -1，表示該維度是動態的，否則表示該維度是明確的，明確維度的 minDims, optDims, maxDims 一定是一樣的。

ONNX文件結構及其增刪改

ONNX 的本質是一個 protobuf 文件
protobuf 則通過 onnx-ml.proto 編譯得到 onnx-ml.pb.h 和 onnx-ml.pb.cc 或 onnx_ml_pb2.py 。
然后用 onnx-ml.pb.cc 和代碼來操作 ONNX 模型文件，實現增刪改。
onnx-ml.proto 則是描述 ONNX 文件是如何組成的，具有什么結構，他是操作 ONNX 經常參照的東西。 https://github.com/onnx/onnx/blob/main/onnx/onnx-ml.proto
ONNX 模型一般是通過常用的模型訓練框架（如 PyTorch 等）導出，當然也可以自己手動構建 ONNX 模型或節點

在這里插入圖片描述

編譯onnx-ml.proto文件

#!/bin/bash# 請修改protoc為你要使用的版本protoc
export LD_LIBRARY_PATH=${@NVLIB64}
protoc=${@PROTOC_PATH}rm -rf pbout
mkdir -p pbout$protoc onnx-ml.proto --cpp_out=pbout --python_out=pbout

ONNX 文件結構

下面一段代碼是 onnx-ml.proto 文件中的一部分關鍵代碼：

message NodeProto {repeated string input = 1;    // namespace Valuerepeated string output = 2;   // namespace Value// An optional identifier for this node in a graph.// This field MAY be absent in ths version of the IR.optional string name = 3;     // namespace Node// The symbolic identifier of the Operator to execute.optional string op_type = 4;  // namespace Operator// The domain of the OperatorSet that specifies the operator named by op_type.optional string domain = 7;   // namespace Domain// Additional named attributes.repeated AttributeProto attribute = 5;// A human-readable documentation for this node. Markdown is allowed.optional string doc_string = 6;
}

表示 onnx 中有節點類型叫 node

它有 input 屬性，是 repeated ，即重復類型，數組
它有 output 屬性，是 repeated ，即重復類型，數組
它有 name 屬性，是 string 類型
后面的數字是 id，一般不用管

關鍵要看的兩個：

repeated ：表示是數組
optional：可選，通常無視即可

我們只關心是否是數組，類型是什么

在這里插入圖片描述

上圖是 onnx 文件結構的一個示意圖，

model：表示整個 ONNX 模型，包含圖結構和解析器格式、opset 版本、導出程序類型等
model.graph：表示圖結構，通常是我們 netron 可視化看到的主要結構
model.graph.node：表示圖中的所有節點，是數組，例如 conv、bn 等算子就是在這里的，通過 input、output 表示節點間的連接關系
model.graph.initializer：權重類的數據大都儲存在這里
model.graph.input：整個模型的輸入儲存在這里，表明哪個節點是輸入節點，shape 是多少
model.graph.output：整個模型的輸出儲存在這里，表明哪個節點是輸入節點，shape 是多少
對于 anchor grid 類的常量數據，通常會存儲在 model.graph.node 中，并指定類型為 Constant 該類型節點在 netron 中可視化時不會顯示出來

讀寫改ONNX文件

PyTorch導出onnx文件

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.onnx
import osclass Model(torch.nn.Module):def __init__(self):super().__init__()self.conv = nn.Conv2d(1, 1, 3, padding=1)self.relu = nn.ReLU()self.conv.weight.data.fill_(1)self.conv.bias.data.fill_(0)def forward(self, x):x = self.conv(x)x = self.relu(x)return x# 這個包對應opset11的導出代碼，如果想修改導出的細節，可以在這里修改代碼
# import torch.onnx.symbolic_opset11
print("對應opset文件夾代碼在這里：", os.path.dirname(torch.onnx.__file__))model = Model()
dummy = torch.zeros(1, 1, 3, 3)
torch.onnx.export(model, # 這里的args，是指輸入給model的參數，需要傳遞tuple，因此用括號(dummy,), # 儲存的文件路徑"demo.onnx", # 打印詳細信息verbose=True, # 為輸入和輸出節點指定名稱，方便后面查看或者操作input_names=["image"], output_names=["output"], # 這里的opset，指，各類算子以何種方式導出，對應于symbolic_opset11opset_version=11, # 表示他有batch、height、width3個維度是動態的，在onnx中給其賦值為-1# 通常，我們只設置batch為動態，其他的避免動態dynamic_axes={"image": {0: "batch", 2: "height", 3: "width"},"output": {0: "batch", 2: "height", 3: "width"},}
)print("Done.!")

創建onnx文件

直接從構建onnx，不經過任何框架的轉換。通過import onnx和onnx.helper提供的make_node，make_graph，make_tensor等等接口我們可以輕易的完成一個ONNX模型的構建。

需要指定 node，initializer，input，output，graph，model 參數

import onnx # pip install onnx>=1.10.2
import onnx.helper as helper
import numpy as np# https://github.com/onnx/onnx/blob/v1.2.1/onnx/onnx-ml.protonodes = [helper.make_node(name="Conv_0",   # 節點名字，不要和op_type搞混了op_type="Conv",  # 節點的算子類型, 比如'Conv'、'Relu'、'Add'這類，詳細可以參考onnx給出的算子列表inputs=["image", "conv.weight", "conv.bias"],  # 各個輸入的名字，結點的輸入包含：輸入和算子的權重。必有輸入X和權重W，偏置B可以作為可選。outputs=["3"],  pads=[1, 1, 1, 1], # 其他字符串為節點的屬性，attributes在官網被明確的給出了，標注了default的屬性具備默認值。group=1,dilations=[1, 1],kernel_shape=[3, 3],strides=[1, 1]),helper.make_node(name="ReLU_1",op_type="Relu",inputs=["3"],outputs=["output"])
]initializer = [helper.make_tensor(name="conv.weight",data_type=helper.TensorProto.DataType.FLOAT,dims=[1, 1, 3, 3],vals=np.array([1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], dtype=np.float32).tobytes(),raw=True),helper.make_tensor(name="conv.bias",data_type=helper.TensorProto.DataType.FLOAT,dims=[1],vals=np.array([0.0], dtype=np.float32).tobytes(),raw=True)
]inputs = [helper.make_value_info(name="image",type_proto=helper.make_tensor_type_proto(elem_type=helper.TensorProto.DataType.FLOAT,shape=["batch", 1, 3, 3]))
]outputs = [helper.make_value_info(name="output",type_proto=helper.make_tensor_type_proto(elem_type=helper.TensorProto.DataType.FLOAT,shape=["batch", 1, 3, 3]))
]graph = helper.make_graph(name="mymodel",inputs=inputs,outputs=outputs,nodes=nodes,initializer=initializer
)# 如果名字不是ai.onnx，netron解析就不是太一樣了
opset = [helper.make_operatorsetid("ai.onnx", 11)
]# producer主要是保持和pytorch一致
model = helper.make_model(graph, opset_imports=opset, producer_name="pytorch", producer_version="1.9")
onnx.save_model(model, "my.onnx")print(model)
print("Done.!")

讀onnx文件

通過graph可以訪問參數，數據是以protobuf的格式存儲的，因此當中的數值會以bytes的類型保存。需要用np.frombuffer方法還原成類型為float32的ndarray。注意還原出來的ndarray是只讀的。

import onnx
import onnx.helper as helper
import numpy as npmodel = onnx.load("demo.change.onnx")#打印信息
print("==============node信息")
# print(helper.printable_graph(model.graph))
print(model)conv_weight = model.graph.initializer[0]
conv_bias = model.graph.initializer[1]# 數據是以protobuf的格式存儲的，因此當中的數值會以bytes的類型保存，通過np.frombuffer方法還原成類型為float32的ndarray
print(f"===================={conv_weight.name}==========================")
print(conv_weight.name, np.frombuffer(conv_weight.raw_data, dtype=np.float32))print(f"===================={conv_bias.name}==========================")
print(conv_bias.name, np.frombuffer(conv_bias.raw_data, dtype=np.float32))

將得到類似如下輸出：

==============node信息
ir_version: 6
producer_name: "pytorch"
producer_version: "1.9"
graph {node {input: "image"input: "conv.weight"input: "conv.bias"output: "3"name: "Conv_0"op_type: "Conv"attribute {name: "dilations"ints: 1ints: 1type: INTS}attribute {name: "group"i: 1type: INT}attribute {name: "kernel_shape"ints: 3ints: 3type: INTS}attribute {name: "pads"ints: 1ints: 1ints: 1ints: 1type: INTS}attribute {name: "strides"ints: 1ints: 1type: INTS}}node {input: "3"output: "output"name: "Relu_1"op_type: "Relu"}name: "torch-jit-export"initializer {dims: 1dims: 1dims: 3dims: 3data_type: 1name: "conv.weight"raw_data: "\000\000\000\000\000\000\200?\000\000\000@\000\000@@\000\000\200@\000\000\240@\000\000\300@\000\000\340@\000\000\000A"}initializer {dims: 1data_type: 1name: "conv.bias"raw_data: "\000\000\000\000"}input {name: "image"type {tensor_type {elem_type: 1shape {dim {dim_param: "batch"}dim {dim_value: 1}dim {dim_param: "height"}dim {dim_param: "width"}}}}}output {name: "output"type {tensor_type {elem_type: 1shape {dim {dim_param: "batch"}dim {dim_value: 1}dim {dim_param: "height"}dim {dim_param: "width"}}}}}
}
opset_import {version: 11
}====================conv.weight==========================
conv.weight [0. 1. 2. 3. 4. 5. 6. 7. 8.]
====================conv.bias==========================
conv.bias [0.]

改onnx文件

由于protobuf任何支持的語言，我們可以使用 c/c++/python/java/c# 等等實現對onnx文件的讀寫操作

掌握onnx和helper實現對onnx文件的各種編輯和修改

增

一般伴隨增加 node 和 tensor

graph.initializer.append(xxx_tensor)
graph.node.insert(0, xxx_node)

例子：比如我們想要在 yolov5s.onnx 的模型前面添加一個預處理，將預處理功能繼承到 ONNX 模型里面，將 opencv 讀到的圖片先預處理成我們想要的格式。這個過程中直接在 ONNX 模型中添加是比較麻煩的，我們的思路是想用 PyTorch 寫一個預處理模塊并導出為 preprocess.onnx ，再將其添加到 yolov5s.onnx 前面。后處理等其他添加節點的操作類似。

步驟：

先用 PyTorch 實現預處理并導出 ONNX 模型 preprocess_onnx
將 preprocess_onnx 中所有節點以及輸入輸出名稱都加上前綴，避免與原模型的名稱沖突
將 yolov5s 中以 image 為輸入的節點，修改為 preprocess_onnx 的輸出節點
將 preprocess_onnx 的 node 全部放到 yolov5s 的 node 中
將 preprocess_onnx 的輸入名稱作為 yolov5s 的 input 名稱

代碼如下：

import torch
import onnx
import onnx.helper as helper# 步驟1
class PreProcess(torch.nn.Module):def __init__(self):super().__init__()self.mean = torch.randn(1, 1, 1, 3) # 這里的均值標準差就直接隨機了，實際模型按需調整self.std = torch.randn(1, 1, 1, 3)def forward(self, x):# 輸入: B H W C uint8# 輸出: B C H W float32, 減255, 減均值除標準差x = x.float()x = (x / 255.0 - self.mean) / self.stdx = x.permute(0, 2, 3, 1)return xpreprocess = PreProcess()
torch.onnx.export(preprocess, (torch.zeros(1, 640, 640, 3, dtype=torch.uint8), ) 'preprocess.onxx')
)preprocess_onnx = onnx.load('preprocess.onnx')
model = onnx.load('yolov5s.onnx')# 步驟2
for item in preprocess_onnx.graph.node:item.name = f"pre/{item.name"for i in range(len(item.input)):item.input[i] = f"pre/{item.input[i]}"for i in range(len(item.output)):item.output[i] = f"pre/{item.output[i]}"# 步驟3
for item in model.graph.node:if item.name == 'Conv_0':item.input[0] = f"pre/{preprocess_onnx.graph.output[0].name}"# 步驟4
for item in pre_onnx.graph.node:model.graph.node.append(item)# 步驟5
input_name = f"pre/{preprocess_onnx.graph.input[0].name}"
model.graph.input[0].CopyFrom(preprocess_onnx.graph.input[0])
model.graph.input[0].name = input_nameonnx.save(model, "yolov5s_with_proprecess.onnx")

刪

刪除節點時需要注意的是要將前一個節點的輸出接到下一個節點的輸入上，就像刪除鏈表節點一樣。

# 刪除一個節點
import onnxmodel = onnx.load('yolox_s.onnx')find_node_with_input = lambda name: [item for item in model.graph.node if name in item.input][0]
find_node_with_output = lambda name: [item for item in model.graph.node if name in item.output][0]remove_nodes = []
for item in model.graph.node:if item.name == "Transpose_236":# 上一個節點的輸出是當前節點的輸入_prev = find_node_with_output(item.input[0])# 下一個節點的輸入是當前節點的輸出_next = find_node_with_input(item.)_next.input[0] = _prev.output[0]remove_nodes.append(item)for item in remove_nodes[::-1]:model.graph.node.remove(item)

改

# 改數據
input_node.name = 'data'# 改掉整個節點
new_item = helper.make_node(...)
item.CopyFrom(new_item)	# `CopyFrom` 是 protobuf 中的函數

通常也可以通過 for loop 去找我們想要的 initializer 或 node 來查看或修改：

for item in model.graph.initializer:if item.name == 'conv1.weight':# do somethingpassfor item in model.graph.node:if item.name == 'Constant':# do somethingpass

例子，修改 yolov5s.onnx 的動態 batch size 靜態尺寸改為靜態 batch size，動態尺寸：

import onnx
import onnx.helper as helpermodel = onnx.load('yolox_s.onnx')
static_batch_size = 4# 原來的輸入尺寸是: [batch, 3, 640, 640], 動態batch size, 固定圖像尺寸
new_input = helper.make_tensor_value_info("images", 1, [static_batch_size, 3, 'height', 'width'])
model.graph.input[0].CopyFrom(new_input)# 原來的輸出尺寸是: [batch, 25200, 85], 動態batch size, 固定錨框數和類別數
new_output = helper.make_tensor_value_info("output", 1, [static_batch_size, 'anchors', 'classes'])
model.graph.output.CopyFrom(new_output)onnx.save(model, 'static_bs4.onnx')