onnx實戰(zhàn)一: 解析yolov5 gpu的onnx優(yōu)化案例:
這是一個英偉達的倉庫, 這個倉庫的做法就是通過用gs對onnx進行修改減少算子然后最后使用TensorRT插件實現(xiàn)算子， 左邊是優(yōu)化過的, 右邊是原版的。 通過這個案例理解原版的onnx的導出流程然后我們看英偉達是怎么拿gs來優(yōu)化這個onnx

原版的export_onnx函數(shù)

先看torch.onnx.export函數(shù)的參數(shù)解釋：

1. model: 要導出的PyTorch模型, 在工程中這里輸入的是訓練好的pt文件

2. im: 這里對應torch.onnx.export的args, 這個是用作模型輸入的示例張量。這幫助ONNX確定輸入的形狀和類型。

3. f: 輸出ONNX模型的文件名或文件對象, 用來指定導出模型的路徑和文件名。

4. verbose (默認為 False): 如果設置為 True，則會打印出模型導出時的詳細日志。

5. opset_version: 導出的ONNX模型的操作集版本。不同的版本可能支持不同的操作。

6. training:

   • torch.onnx.TrainingMode.TRAINING: 表示模型處于訓練模式。
   • torch.onnx.TrainingMode.EVAL: 表示模型處于評估模式。

7. do_constant_folding (默認為 True): 當設置為 True，導出過程中會嘗試簡化模型，將常量子圖折疊為一個常量節(jié)點。

8. input_names: 為模型的輸入提供名稱, 參數(shù)規(guī)定是數(shù)組

9. output_names: 為模型的輸出提供名稱, 參數(shù)規(guī)定是數(shù)組

10. dynamic_axes: 為模型的輸入/輸出定義動態(tài)軸。對于那些維度在推理時可能會發(fā)生變化的情況（例如，批處理大小），此參數(shù)允許指定哪些軸是動態(tài)的。這里images是輸入, 本來是1x3x640x640, 這里通過指定把0, 2, 3維度變成了動態(tài)軸的輸入, 第二個維度是3這個還是固定的。如果使用動態(tài), 可以輸入任意數(shù)量和任意大小的圖片而不是規(guī)定的單張640x640

• 'images': 對應的張量名稱。
  • 0: 'batch': 表示第0個維度（即批處理維度）是動態(tài)的，并命名為’batch’。
  • 2: 'height': 表示第2個維度（即圖像的高度）是動態(tài)的。
  • 3: 'width': 表示第3個維度（即圖像的寬度）是動態(tài)的。
• 'output': 對應的張量名稱。
  • 0: 'batch': 表示第0個維度（即批處理維度）是動態(tài)的。
  • 1: 'anchors': 表示第1個維度是動態(tài)的。

11. dynamic (沒有在給定的函數(shù)調(diào)用中明確給出，但可以從上下文推斷):

• True: 如果你想讓某些軸動態(tài)，你可以設置此參數(shù)為True。
• False: 表示不使用動態(tài)軸。

導出了onnx之后開始做onnxsim

1. model_onnx, check = onnxsim.simplify(...):使用onnxsim的simplify方法簡化模型。它返回簡化后的onnx模型和一個布爾值check，表示簡化是否成功。

2. 在對動態(tài)輸入的onnx導出的時候, dynamic_input_shape=dynamic是不夠的，還要把輸入給他，讓onnxsim更加謹慎的優(yōu)化onnx, 確保滿足我們給他的輸出，所以這里多了一個input_shapes={'images': list(im.shape)} if dynamic else None

def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorstr('ONNX:')):# YOLOv5 ONNX exporttry:check_requirements(('onnx',))import onnxLOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...')f = file.with_suffix('.onnx')torch.onnx.export(model,im,f,verbose=False,opset_version=opset,training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,do_constant_folding=not train,input_names=['images'],output_names=['output'],dynamic_axes={'images': {0: 'batch',2: 'height',3: 'width'},  # shape(1,3,640,640)'output': {0: 'batch',1: 'anchors'}  # shape(1,25200,85)} if dynamic else None)# Checksmodel_onnx = onnx.load(f)  # load onnx modelonnx.checker.check_model(model_onnx)  # check onnx model# Metadatad = {'stride': int(max(model.stride)), 'names': model.names}for k, v in d.items():meta = model_onnx.metadata_props.add()meta.key, meta.value = k, str(v)onnx.save(model_onnx, f)# Simplifyif simplify:try:check_requirements(('onnx-simplifier',))import onnxsimLOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')model_onnx, check = onnxsim.simplify(model_onnx,dynamic_input_shape=dynamic,input_shapes={'images': list(im.shape)} if dynamic else None)assert check, 'assert check failed'onnx.save(model_onnx, f)except Exception as e:LOGGER.info(f'{prefix} simplifier failure: {e}')LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')return fexcept Exception as e:LOGGER.info(f'{prefix} export failure: {e}')

更改過的export_onnx函數(shù)

1. 首先是把onnx的輸出由一個改成了3個, 然后指定動態(tài)輸出, 因為有多個輸出，全部都把他們的batch, width, height指定為動態(tài)的，滿足不同的輸入輸出。 不過這邊的問題是看起來是只改了最后的輸出，但是前面在yolo.py的地方已經(jīng)把sigmoid后面的計算都干掉了, 因為后面的計算映射了一堆的算子導致了在計算圖太冗余

這一坨全部不要了就保留sigmoid就可以了，然后就是直接硬編碼t就是int32

diff --git a/models/yolo.py b/models/yolo.py
index 02660e6..c810745 100644
--- a/models/yolo.py
+++ b/models/yolo.py
@@ -55,29 +55,15 @@ class Detect(nn.Module):z = []  # inference outputfor i in range(self.nl):x[i] = self.m[i](x[i])  # conv
-            bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
-            x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
-
-            if not self.training:  # inference
-                if self.onnx_dynamic or self.grid[i].shape[2:4] != x[i].shape[2:4]:
-                    self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)
-
-                y = x[i].sigmoid()
-                if self.inplace:
-                    y[..., 0:2] = (y[..., 0:2] * 2 + self.grid[i]) * self.stride[i]  # xy
-                    y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
-                else:  # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
-                    xy, wh, conf = y.split((2, 2, self.nc + 1), 4)  # y.tensor_split((2, 4, 5), 4)  # torch 1.8.0
-                    xy = (xy * 2 + self.grid[i]) * self.stride[i]  # xy
-                    wh = (wh * 2) ** 2 * self.anchor_grid[i]  # wh
-                    y = torch.cat((xy, wh, conf), 4)
-                z.append(y.view(bs, -1, self.no))
-
-        return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(z, 1), x)
+            y = x[i].sigmoid()
+            z.append(y)
+        return zdef _make_grid(self, nx=20, ny=20, i=0):d = self.anchors[i].device
-        t = self.anchors[i].dtype
+        # t = self.anchors[i].dtype
+        # TODO(tylerz) hard-code data type to int
+        t = torch.int32shape = 1, self.na, ny, nx, 2  # grid shapey, x = torch.arange(ny, device=d, dtype=t), torch.arange(nx, device=d, dtype=t)if check_version(torch.__version__, '1.10.0'):  # torch>=1.10.0 meshgrid workaround for torch>=0.7 compatibility
-- 
2.36.0

2. onnxsim這里跟之前是一樣的, 也是直接onnxsim, 如果動態(tài)的要給輸出給onnxsim然后讓它更加的謹慎，滿足需求

3. 這后面的重點是增加了用onnx-surgon來更改onnx, 先把整個onnx導入進來，然后使用然后用Variable做模型的輸出, 這里做四個模型輸出, 分別是DecodeNumDetection, DecodeDetectionBoxes, DecodeDetectionScores, DecodeDetectionClasses

4. 然后設置一個attrs, gs設置的attrs使用字典的格式弄的。這里設置max_stride， num_classes, anchors, prenms_score_threshold四個屬性，這些屬性的操作會在TensorRT中實現(xiàn)的

5. decode_plugin是中間的節(jié)點，這個節(jié)點上面是inputs, 下面是四個不同的decodes, 這里就是把這個nodes做出來了

6. 然后在整體的網(wǎng)絡上添加這個節(jié)點，然后再把輸出改成這個節(jié)點的輸出保持一致，在計算圖中把其他的節(jié)點claenup()清潔掉, 最后導出

def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorstr('ONNX:')):# YOLOv5 ONNX export# try:check_requirements(('onnx',))import onnxLOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...')f = file.with_suffix('.onnx')print(train)torch.onnx.export(model,im,f,verbose=False,opset_version=opset,training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,do_constant_folding=not train,input_names=['images'],output_names=['p3', 'p4', 'p5'],dynamic_axes={'images': {0: 'batch',2: 'height',3: 'width'},  # shape(1,3,640,640)'p3': {0: 'batch',2: 'height',3: 'width'},  # shape(1,25200,4)'p4': {0: 'batch',2: 'height',3: 'width'},'p5': {0: 'batch',2: 'height',3: 'width'}} if dynamic else None)# Checksmodel_onnx = onnx.load(f)  # load onnx modelonnx.checker.check_model(model_onnx)  # check onnx model# Simplifyif simplify:# try:check_requirements(('onnx-simplifier',))import onnxsimLOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')model_onnx, check = onnxsim.simplify(model_onnx,dynamic_input_shape=dynamic,input_shapes={'images': list(im.shape)} if dynamic else None)assert check, 'assert check failed'onnx.save(model_onnx, f)# except Exception as e:#     LOGGER.info(f'{prefix} simplifier failure: {e}')# add yolov5_decoding:import onnx_graphsurgeon as onnx_gsimport numpy as npyolo_graph = onnx_gs.import_onnx(model_onnx)p3 = yolo_graph.outputs[0]p4 = yolo_graph.outputs[1]p5 = yolo_graph.outputs[2]decode_out_0 = onnx_gs.Variable("DecodeNumDetection",dtype=np.int32)decode_out_1 = onnx_gs.Variable("DecodeDetectionBoxes",dtype=np.float32)decode_out_2 = onnx_gs.Variable("DecodeDetectionScores",dtype=np.float32)decode_out_3 = onnx_gs.Variable("DecodeDetectionClasses",dtype=np.int32)decode_attrs = dict()decode_attrs["max_stride"] = int(max(model.stride))decode_attrs["num_classes"] = model.model[-1].ncdecode_attrs["anchors"] = [float(v) for v in [10,13, 16,30, 33,23, 30,61, 62,45, 59,119, 116,90, 156,198, 373,326]]decode_attrs["prenms_score_threshold"] = 0.25decode_plugin = onnx_gs.Node(op="YoloLayer_TRT",name="YoloLayer",inputs=[p3, p4, p5],outputs=[decode_out_0, decode_out_1, decode_out_2, decode_out_3],attrs=decode_attrs)yolo_graph.nodes.append(decode_plugin)yolo_graph.outputs = decode_plugin.outputsyolo_graph.cleanup().toposort()model_onnx = onnx_gs.export_onnx(yolo_graph)d = {'stride': int(max(model.stride)), 'names': model.names}for k, v in d.items():meta = model_onnx.metadata_props.add()meta.key, meta.value = k, str(v)onnx.save(model_onnx, f)LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')return f# except Exception as e:#     LOGGER.info(f'{prefix} export failure: {e}')