将yolov5s部署到安卓上实战经验总结

最近需要在手机端实现一个目标检测的功能，于是选择了小巧又在目标检测方面表现很好的yolov5s，官网下载yolov5代码，用自己做的数据集进行了训练，然后把模型转换成torchscript格式，这些过程网上都有很多讲解，不再赘述。主要讲一下在安卓上推理的代码。

pytorch在安卓上的使用官方demo，主要代码如下：

    Bitmap bitmap = null;
    Module module = null;
    try {
      // creating bitmap from packaged into app android asset 'image.jpg',
      // app/src/main/assets/image.jpg
      bitmap = BitmapFactory.decodeStream(getAssets().open("image.jpg"));
      // loading serialized torchscript module from packaged into app android asset model.pt,
      // app/src/model/assets/model.pt
      module = LiteModuleLoader.load(assetFilePath(this, "model.pt"));
    } catch (IOException e) {
      Log.e("PytorchHelloWorld", "Error reading assets", e);
      finish();
    }

    // showing image on UI
    ImageView imageView = findViewById(R.id.image);
    imageView.setImageBitmap(bitmap);

    // preparing input tensor
    final Tensor inputTensor = TensorImageUtils.bitmapToFloat32Tensor(bitmap,
        TensorImageUtils.TORCHVISION_NORM_MEAN_RGB, TensorImageUtils.TORCHVISION_NORM_STD_RGB, MemoryFormat.CHANNELS_LAST);

    // running the model
    final Tensor outputTensor = module.forward(IValue.from(inputTensor)).toTensor();

    // getting tensor content as java array of floats
    final float[] scores = outputTensor.getDataAsFloatArray();

    // searching for the index with maximum score
    float maxScore = -Float.MAX_VALUE;
    int maxScoreIdx = -1;
    for (int i = 0; i < scores.length; i++) {
      if (scores[i] > maxScore) {
        maxScore = scores[i];
        maxScoreIdx = i;
      }
    }

    String className = ImageNetClasses.IMAGENET_CLASSES[maxScoreIdx];

但是这段代码中用的模型不是yolov5，直接用于yolov5的模型是跑不通的，首先计算outputTensor的时候直接把模型输出toTensor()，这个会报错，报错讲说期望Tensor类型但是给了个Tuple，由此可知模型的输出IValue其内置类型是Tuple，于是toTuple然后取第一个元素再toTensor()就可以了。原因是yolov5的输出在Tensor外面又包装了一层，组成了一个Tuple。

然后是结果scores的解析方法，对于yolov5，当有n个目标类别的时候，这个scores的含义是[x,y,w,h,conf,type1score,type2score,......typenscore,x,y,w,h,conf,type1score,type2score,....typenscore......]，一直重复25200次，其中x,y是目标框的中心坐标，w,h是目标框的宽高，conf是框的置信度，后面分别是n个类别的得分。所以自然不能用上述代码中的方法取结果。

等我修改完这两处之后，代码可以正常运行，但奇怪的是在python上运行训练好的模型，结果是非常好的，基本95%的时候都可以获取到目标物体在图像中的最小外接矩形，其它5%也只是偏移一点点，但到了手机上，这个结果常常不准确，检测框没有包住目标物体的所有部分是很大概率的事，一开始我怀疑是模型转换的时候丢失了精度，但后来发现转换成torchscript并没有量化，并且在不量化的情况下，模型没必要把一些参数进行修改，这不是努力降精度吗？不合常理。于是仔细看了下yolov5源码中的推理部分，发现图片在进入模型之前，进行了/255的归一化操作。于是乎问题聚集到了原来代码中的这一行：

TensorImageUtils.bitmapToFloat32Tensor(bitmap,
TensorImageUtils.TORCHVISION_NORM_MEAN_RGB, TensorImageUtils.TORCHVISION_NORM_STD_RGB, MemoryFormat.CHANNELS_LAST);
经过了多次调试，终于发现这个函数其实是对bitmap的像素值进行了/255的归一化后，再使用传入的均值数组和标准差数组对归一化过的数值进行了Z-score归一化。Z-score归一化的目的原本是为了让数据符合标准正态分布，但是进入TensorImageUtils类可以看到：

public static float[] TORCHVISION_NORM_MEAN_RGB = new float[]{0.485F, 0.456F, 0.406F};
public static float[] TORCHVISION_NORM_STD_RGB = new float[]{0.229F, 0.224F, 0.225F};

即使用了事先固定的均值和标准差，而不是传入数据的均值和标准差，所以不一定可以得到符合标准正态分布的数据。但是这不重要，因为我要的是直接不作Z-score归一化，只/255就可以了，于是我自定义了一个值为0的均值数组，和值为1的标准差数组，然后传入这个函数，就保证了结果相当于没有做Z-score归一化。原因是Z-score归一化公式如下：

x* = ( x − μ ) / σ

我的最终关键代码如下：注意处理结果的部分，因为我是图片中一定只有0或1个目标检测框，所以我没有使用NMS（非极大值抑制）算法。如果你的图片中有多个检测框，则必须用NMS。我只有两个类别，所以idcnt计算是score.length/7，也就是score.length/(4+1+类别数)。


        model = Module.load(path);

        float[] TORCHVISION_NORM_MEAN_RGB = new float[]{0F, 0F, 0F};
        float[] TORCHVISION_NORM_STD_RGB = new float[]{1F, 1F, 1F};
        Tensor inputTensor = TensorImageUtils.bitmapToFloat32Tensor(newBitmap, TORCHVISION_NORM_MEAN_RGB, TORCHVISION_NORM_STD_RGB);
        // running the model
        IValue value = IValue.from(inputTensor);

        Tensor outputTensor_ori = model.forward(value).toTuple()[0].toTensor();
        // getting tensor content as java array of floats
        float[] scores = outputTensor_ori.getDataAsFloatArray();
        // searching for the index with maximum score
        float maxScore = 0.85F;
        int maxScoreIdx = -1;
        int idcnt = scores.length / 7;
        for (int i = 0; i < idcnt; i++) {
            int exist = i*7+4;
            int j = exist+1+type;
            if (scores[exist] > 0.25F && scores[j] > maxScore) {
                maxScore = scores[j];
                maxScoreIdx = i;
            }
        }
        if (maxScoreIdx == -1) {
            return false;
        }

        float tx = scores[maxScoreIdx*7];
        float ty = scores[maxScoreIdx*7+1];
        float tw = scores[maxScoreIdx*7+2];
        float th = scores[maxScoreIdx*7+3];

        float ltx = (tx-tw/2);
        float lty = (ty-th/2);
        float rbx = (tx+tw/2);
        float rby = (ty+th/2);
        drawROI(newBitmap, (int)ltx, (int)lty, (int)rbx, (int)rby);