在Unity上使用Sentis部署yolov8与添加NMS，并编译到手机上_barracuda　vs　sentis

项目地址

项目：https://github.com/212534/Unity-Sentis-YOLOv8

Demo apk：链接：https://pan.baidu.com/s/1agTZRhnCzgT5P5HtuUvgWQ?pwd=ydj7
提取码：ydj7
–来自百度网盘超级会员V5的分享

效果展示

这是在电脑上的测试，用的摄像头拍屏幕

安装Sentis

可以把Sentis看作Barracuda的升级版。
在Package里装com.unity.sentis

使用Sentis推理yolov8的onnx

这一部分比较简单，教程非常多

using System.Collections;
using System.Collections.Generic;
using Unity.Sentis;
using UnityEngine;
using Unity.Sentis.Layers;

public class Test: MonoBehaviour
{
    public ModelAsset modelAsset;
    private Model model;
    private IWorker worker;
    Ops ops;

    void Start()
    {
        model = ModelLoader.Load(modelAsset);
        worker = WorkerFactory.CreateWorker(BackendType.GPUCompute, model);
        ops = WorkerFactory.CreateOps(BackendType.GPUCompute, null);
    }

    public void Predict(WebCamTexture camImage) //我这里使用的是摄像头图像，你也可以用普通图片。
    {
        using Tensor inputImage = TextureConverter.ToTensor(camImage,width:640,height:640, channels: 3); //对输入的图像做处理

        var m_Inputs = new Dictionary<string, Tensor> 
        {
            {"images", inputImage }
        };

        worker.Execute(m_Inputs);//执行推理

        var output0 = worker.PeekOutput("output0") as TensorFloat;  //获取输出结果
        output0.MakeReadable(); //从GPU中取出数据，经过这一步之后就可以读取output0中的数据了
    }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35

直接把onnx文件拖上去就行了。

YOLOv8的输出格式

需要注意的是，yolov8的onnx的输出结果是一个1x84x8400的张量（当类别数不同时会有变化），具体含义可以参考这篇文章https://blog.csdn.net/yangkai6121/article/details/133843368
简单来讲，这里的8400是8400个预测框。而每个预测框包含以下信息：84 = 边界框预测4 + 数据集类别80
本文使用的yolov8是在coco数据集上训练的，包含80个类别，所以这里是4+80，具体情况与你的训练用的数据集类别有关。
1x84x8400这个输出很明显没法直接使用。所以需要加一个nms来后处理

NMS

yolo属于单阶段的目标检测算法，因此还要对预测框进行后处理。常见的做法是在Unity里用C#自己实现nms，但nms本身就是一个性能开销很大的后处理，自己用C#实现在性能上不划算。

如果是了解目标检测部署的同学，应该知道onnx在opset11及以上版本支持了内置nms层。那么理论上，只要在yolov8的检测头里把torch.ops.torchvision.nms加进去就可以导出自带nms层的ONNX模型。像下图一样

实际上经过测试，Sentis确实可以识别出onnx的内置nms层，但是会报错。

原因可能在于，Sentis内实现的nms层与torch.ops.torchvision.nms在input的shape上存在区别。

举例来讲，在python中对yolov8使用torch.ops.torchvision.nms时，它要求的boxes输入shape为8400x4，scores为8400.
但是Sentis实现的nms层要求的shape分别为，1x8400x4与1x1x8400。

使用Sentis给模型加入NMS层

既然torch.ops.torchvision.nms不能用，那我们可以在Unity里用代码把Sentis实现的nms层加到模型里。
先看一下Sentis中NonMaxSuppression()的输入参数。详细的可以看这里
https://docs.unity3d.com/Packages/com.unity.sentis@1.2/api/Unity.Sentis.Layers.NonMaxSuppression.html

大体来讲与torch.ops.torchvision.nms的输入参数比较一致。
我们需要解决的问题便是把yolov8输出的1x84x8400张量拆解为boxes与scores，也就是NonMaxSuppression()输入。而Sentis是支持在Unity中对模型结构进行修改的。

1.首先我们给模型加入三个新的输入，这也是NonMaxSuppression()所需要的输入。使用新增输入的形式设置这三个值，可以在程序运行时对它们动态修改。

		//Set input
        model.AddInput("maxOutputBoxesPerClass",DataType.Int, new SymbolicTensorShape(1)); //每个类别最多返回的边框数量。
        model.AddInput("iouThreshold",DataType.Float, new SymbolicTensorShape(1)); //iou阈值
        model.AddInput("scoreThreshold",DataType.Float, new SymbolicTensorShape(1)); //置信度阈值
1
2
3
4

2.对1x8400x4进行拆分。比较反人类的是，Sentis在添加层时，有些层的初始化只能输入字符串作为参数，这些字符串实际上是网络里的输出名，输入名，常量名等。所以这里要给模型加入几个常量。
我在注释里加入了每一层的shape变化，方便各位理解这个过程。如果你没搞过深度学习，不理解每一层的含义，可以去官方文档里看具体解释。
https://docs.unity3d.com/Packages/com.unity.sentis@1.2/api/Unity.Sentis.Layers.html

 		//Set constants   
        model.AddConstant(new Constant("0", new int[] { 0 })); 
        model.AddConstant(new Constant("1", new int[] { 1 }));
        model.AddConstant(new Constant("4", new int[] { 4 }));
        model.AddConstant(new Constant("84", new int[] { 84 })); //这个值根据你自己的类别数进行修改  4 + 类别数

        //Add layers    
        model.AddLayer(new Slice("boxCoords0", "output0", "0", "4", "1")); //1x84x8400 -> 1x4x8400
        model.AddLayer(new Transpose("boxCoords", "boxCoords0", new int[] { 0, 2, 1 })); // 1x4x8400 -> 1x8400x4  每个检测框的xywh

        model.AddLayer(new Slice("scores0", "output0", "4", "84", "1"));  //1x84x8400 -> 1x80x8400
        model.AddLayer(new ReduceMax("scores", new[] { "scores0", "1" }));  //1x80x8400 -> 1x1x8400   最有可能类别的置信度
        model.AddLayer(new ArgMax("classIDs", "scores0", 1)); //1x80x8400 -> 1x1x8400 最有可能类别的index
1
2
3
4
5
6
7
8
9
10
11
12
13

3.加入NMS层。其中，yolov8输出的边界框信息是xywh的格式，因此设置centerPointBox: CenterPointBox.Center

		model.AddLayer(new NonMaxSuppression("nmsOutput", "boxCoords", "scores",
            "maxOutputBoxesPerClass", "iouThreshold", "scoreThreshold",
            centerPointBox: CenterPointBox.Center
        ));
1
2
3
4

4.给网络加入新的输出

		model.AddOutput("boxCoords"); 
        model.AddOutput("classIDs");
        model.AddOutput("nmsOutput");  
1
2
3

最终变成这样：

	void Start()
    {
        model = ModelLoader.Load(modelAsset);
        worker = WorkerFactory.CreateWorker(BackendType.GPUCompute, model);
        ops = WorkerFactory.CreateOps(BackendType.GPUCompute, null);
        
        //Set input
        model.AddInput("maxOutputBoxesPerClass",DataType.Int, new SymbolicTensorShape(1));
        model.AddInput("iouThreshold",DataType.Float, new SymbolicTensorShape(1));
        model.AddInput("scoreThreshold",DataType.Float, new SymbolicTensorShape(1));

        //Set constants   
        model.AddConstant(new Constant("0", new int[] { 0 }));
        model.AddConstant(new Constant("1", new int[] { 1 }));
        model.AddConstant(new Constant("4", new int[] { 4 }));
        model.AddConstant(new Constant("84", new int[] { 84 }));

        //Add layers    
        model.AddLayer(new Slice("boxCoords0", "output0", "0", "4", "1")); //1x84x8400 -> 1x4x8400
        model.AddLayer(new Transpose("boxCoords", "boxCoords0", new int[] { 0, 2, 1 })); // 1x4x8400 -> 1x8400x4

        model.AddLayer(new Slice("scores0", "output0", "4", "84", "1"));  //1x84x8400 -> 1x80x8400
        model.AddLayer(new ReduceMax("scores", new[] { "scores0", "1" }));  //1x80x8400 -> 1x1x8400   最有可能类别的置信度
        model.AddLayer(new ArgMax("classIDs", "scores0", 1)); //1x80x8400 -> 1x1x8400 最有可能类别的index

        model.AddLayer(new NonMaxSuppression("nmsOutput", "boxCoords", "scores",
            "maxOutputBoxesPerClass", "iouThreshold", "scoreThreshold",
            centerPointBox: CenterPointBox.Center
        ));

        model.AddOutput("boxCoords");
        model.AddOutput("classIDs");
        model.AddOutput("nmsOutput");   
    }
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

使用NMS的输出结果

NonMaxSuppression()的输出是一个Nx3的数据。
N指的是最终保留了几个框，而3则代表[boxID,clsID,boxCoordID] ，有了clsID以及boxCoordID便可以从模型输出的boxCoords以及classIDs中提取出最终预测框的xywh和类别。
1.修改一下网络的输入，因为前文中给网络新加了三个输入，因此这里需要改变输入的数据。

		using Tensor inputImage = TextureConverter.ToTensor(camImage,width:640,height:640, channels: 3);
        var m_Inputs = new Dictionary<string, Tensor>
        {
            {"images", inputImage },
            {"maxOutputBoxesPerClass", new TensorInt(new TensorShape(1), new int[] { maxOutputBoxesPerClass })},
            {"iouThreshold", new TensorFloat(new TensorShape(1), new float[] { iouThreshold })},
            {"scoreThreshold",new TensorFloat(new TensorShape(1), new float[] { scoreThreshold })}
        };
        worker.Execute(m_Inputs);
1
2
3
4
5
6
7
8
9

2.获取网络输出

		var boxCoords = worker.PeekOutput("boxCoords") as TensorFloat;  //1x8400x4 所有的预测框的xywh  
        var nmsOutput = worker.PeekOutput("nmsOutput") as TensorInt;  //Nx3, N指的是最终保留了几个框. 返回结果为[boxID,clsID,boxCoordID]   
        var classIDs = worker.PeekOutput("classIDs") as TensorInt; //1x1x8400  每个预测框的类别id
1
2
3

3.处理网络的输出，方便后续使用。
模型的输出结果存在于GPU上，一般来讲，需要先从GPU上取出数据才能用CPU进行处理，这会导致性能的浪费。因此我们可以使用ops调用Sentis的算子，直接在GPU上对网络的结果进行处理，这样可以节省一些性能开销。

		if(nmsOutput.shape[0] == 0)//如果NMS之后一个框也未剩下 则终止
        {
            return;
        }

        var boxCoordIDs = ops.Slice(nmsOutput, new int[] { 2 }, new int[] { 3 }, new int[] { 1 }, new int[] { 1 }); //Nx3 -> Nx1  取出boxCoordID
        var boxCoordIDsFlat = boxCoordIDs.ShallowReshape(new TensorShape(boxCoordIDs.shape.length)) as TensorInt;  //Nx1 -> N  展平
        var output = ops.Gather(boxCoords, boxCoordIDsFlat, 1) as TensorFloat; //1x8400x4 -> 1xNx4
        var labelIDs = ops.Gather(classIDs, boxCoordIDsFlat, 2) as TensorInt; //1x1x8400 -> N

        output.MakeReadable(); //将处理好的结果从GPU取出
        labelIDs.MakeReadable();
1
2
3
4
5
6
7
8
9
10
11
12

4.经过前面处理过的输出，从GPU里取出来就可以直接使用。一个简单示例如下。

		for(int i = 0; i<output.shape[1]; i++)
        {
            float x=output[0,i,0];
            float y=output[0,i,1];
            float w=output[0,i,2];
            float h=output[0,i,3];
            DrawBoundingBoxes(x,y,w,h,labelIDs[i]);//绘制一个Bounding Box
        }
1
2
3
4
5
6
7
8

最终如下所示：

	public void Predict(WebCamTexture camImage)
    {
        using Tensor inputImage = TextureConverter.ToTensor(camImage,width:640,height:640, channels: 3);
        var m_Inputs = new Dictionary<string, Tensor>
        {
            {"images", inputImage },
            {"maxOutputBoxesPerClass", new TensorInt(new TensorShape(1), new int[] { maxOutputBoxesPerClass })},
            {"iouThreshold", new TensorFloat(new TensorShape(1), new float[] { iouThreshold })},
            {"scoreThreshold",new TensorFloat(new TensorShape(1), new float[] { scoreThreshold })}
        };
        worker.Execute(m_Inputs);
        var boxCoords = worker.PeekOutput("boxCoords") as TensorFloat;  //1x8400x4 所有的预测框的xywh
        var nmsOutput = worker.PeekOutput("nmsOutput") as TensorInt;  //Nx3, N指的是最终保留了几个框. 返回结果为[boxID,clsID,boxCoordID]   
        var classIDs = worker.PeekOutput("classIDs") as TensorInt; //1x1x8400  每个预测框的类别id

        if(nmsOutput.shape[0] == 0)//如果NMS之后一个框也未剩下 则终止
        {
            return;
        }

        var boxCoordIDs = ops.Slice(nmsOutput, new int[] { 2 }, new int[] { 3 }, new int[] { 1 }, new int[] { 1 }); //Nx3 -> Nx1  取出boxCoordID
        var boxCoordIDsFlat = boxCoordIDs.ShallowReshape(new TensorShape(boxCoordIDs.shape.length)) as TensorInt;  //Nx1 -> N  展平
        var output = ops.Gather(boxCoords, boxCoordIDsFlat, 1) as TensorFloat; //1x8400x4 -> 1xNx4
        var labelIDs = ops.Gather(classIDs, boxCoordIDsFlat, 2) as TensorInt; //1x1x8400 -> N

        output.MakeReadable();
        labelIDs.MakeReadable();

        for(int i = 0; i<output.shape[1]; i++)
        {
            float x=output[0,i,0];
            float y=output[0,i,1];
            float w=output[0,i,2];
            float h=output[0,i,3];
            DrawBoundingBoxes(x,y,w,h,labelIDs[i]);//绘制一个Bounding Box
        }
    }
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

导出到手机

如果你要打包到手机端的话，记得点一下onnx文件，然后点这里。

下面是手机端的运行效果，我手机的CPU是骁龙850，可以发现帧数非常低。主要是我把推理放到了Update()里，所以性能开销很大，因此还有很多优化空间。