YOLOv5 Android(完结)

最近在搞移动端的模型部署包括coreML IOS端的部署，tf-Lite Android端的部署。但是自己并不是前端开发工程师，所以移动端的代码只是修修改改，主要还是如何将模型转换到对应的格式并可以运行。

这篇主要说下YOLOv5s在Android端的部署，模型是官方最新的的pytorch转ONNX，然后自己再转tf-lite，中间需要自己去修改NMS相关部分的代码，遇到一堆坑。

效果录像

流程概述

1. Step 1：用pytorch训练YoloV5s模型，并生成pt模型文件。
2. Step 2：将pt模型加载并利用pytorch torch.jit.trace功能生成jit模型。
3. Step 3：将jit模型转换（torch.onnx.export）到onnx模型。
4. Step 4：利用onnx-tf脚本将onnx导出到tensorflow支持的savedModel模型。
5. Step 5：tensorflow加载savedModel模型并编写后处理算法并添加到模型尾部然后保存为savedModel模型。
6. Step 6：利用tensorflow-lite将 savedModel模型转换为tflite模型。
7. Step 7：手机端tflite模型部署。

 

Step 1：用pytorch训练YoloV5s模型，并生成pt模型文件。

       工程链接：https://github.com/ultralytics/yolov5

       具体看readme教程

Step 2：将pt模型加载并利用pytorch torch.jit.trace功能生成jit模型。

        这一步主要是转onnx时确保pytorch模型不要有onnx不支持的算子，如果只是卷积操作基本问题不大，但是还是会有奇奇怪怪的问题需要去克服。

Step 3：将jit模型转换（torch.onnx.export）到onnx模型。

这一步主要就是几行代码问题，就用export.py 即可. 需要注意的是

model.model[-1].export = False 这里要改成 False 否则无法导出，主要还是版本问题，tf新版本有些是可以，但是会导致一个AddV2算子无法支持新版本只有tf2.1.1支持。但是tf2.1.1版本也会有其他问题需要改比如NMSV3不支持需要改成V4.以及算法中不要用类似省略号去张量的操作[…,:,:,3]，tf-lite会报错，直接用：，：，：代替。还有tf-lite不支持超过4d维度的张量操作。需要修改yolo.py文件模型代码如下。实际上这里就是模型输出后NMS前导工作。

def forward(self, x):

        # x = x.copy()  # for profiling

        z = []  # inference output

        self.training |= self.export

        for 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)

            if not self.training:  # inference

                x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()

                if self.grid[i].shape[2:4] != x[i].shape[2:4]:

                    self.grid[i] = self._make_grid(nx, ny).to(x[i].device)

                y = x[i].sigmoid()

                y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i]  # xy

                y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh

                z.append(y.view(bs, -1, self.no))

        return x if self.training else (torch.cat(z, 1), x)

 

Step 4：利用onnx-tf脚本将onnx导出到tensorflow支持的savedModel模型。

这个倒是很简单一行命令：onnx模型地址， 导出savedmodel位置

onnx-tf convert -i ./models/5.onnx -o ./model_new

注：如果你用的老版本onnx-tf 导出模型可能是冻结模型pb，我觉得还是用savedmodel最好，冻结模型pb已经用的少了。

Step 5：tensorflow加载savedModel模型并编写NMS算法并添加到模型尾部然后保存为savedModel模型。

这一块主要是先导入savedModel利用TF的tf.function()添加可追踪对象在，主要是添加两部分代码：

图片预处理代码
def preprocess(image):

        image = image[:, :, ::-1]

        image = tf.transpose(image, (2, 0, 1))

        image = tf.cast(image, tf.float32)/255.0

        image = tf.expand_dims(image, axis=0)

        return image

NMS算法的前后代码
def mynms(inp):

        stride = [8., 16., 32.]

        anchor_grid = [[[[[10., 13.]]],

                        [[[16., 30.]]],

                        [[[33., 23.]]]],

                       [[[[30., 61.]]],

                        [[[62., 45.]]],

                        [[[59., 119.]]]],

                       [[[[116., 90.]]],

                        [[[156., 198.]]],

                        [[[373., 326.]]]]]

        z = []

        for i, one in enumerate(inp):

            layer_ = one

            ny, nx = layer_.shape[-2:]

            layer_ = tf.transpose(tf.reshape(layer_, [3, 85, ny, nx]), [0, 2, 3, 1])

            yv, xv = tf.meshgrid(tf.range(0, ny), tf.range(0, nx), indexing='ij')

            l_meshgrid = tf.cast(tf.reshape(tf.stack([xv, yv], 2), [1, ny, nx, 2]), tf.float32)

            layer_ = tf.sigmoid(layer_)

            cxcy = (layer_[:, :, :, 0:2] * 2. - 0.5 + l_meshgrid) * stride[i]

            wh = (layer_[:, :, :, 2:4] * 2) ** 2 * anchor_grid[i]

            layer_ = tf.concat([cxcy, wh, layer_[:, :, :, 4:]], axis=-1)

            layer_ = tf.reshape(layer_, [-1, 85])

            z.append(layer_)

        layers_out = tf.concat(z, axis=0)

        layers_out = tf.expand_dims(layers_out, 0)

        boxes = layers_out[0, :, :4]

        # cx, cy, w, h to x0 y0 x1 y1

        x0 = boxes[:, 0] - boxes[:, 2] / 2

        y0 = boxes[:, 1] - boxes[:, 3] / 2

        x1 = boxes[:, 0] + boxes[:, 2] / 2

        y1 = boxes[:, 1] + boxes[:, 3] / 2

        boxes = tf.stack((x0, y0, x1, y1), 1)

        conf = layers_out[0, :, 4]

        classes = layers_out[0, :, 5:]

        selected_indices, _ = tf.image.non_max_suppression_with_scores(boxes, conf, 30, 0.3)

        boxes = tf.gather(boxes, selected_indices, name="boxes")

        conf = tf.gather(conf, selected_indices, name="conf")

        classes = tf.gather(classes, selected_indices, name="classes")

        return boxes, conf, classes

Step 6：利用tensorflow-lite将 savedModel模型转换为tflite模型。

这里需要注意就是GPU可能会显示不够需要内存自增长设置。

import tensorflow as tf

import numpy as np

import predict_

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')

for gpu in gpus:

    tf.config.experimental.set_memory_growth(gpu, True)

# input_arrays = ["input"]

# output_arrays = ["boxes", 'conf', 'classes']

# converter = tf.lite.TFLiteConverter.from_saved_model("/media/george/DATASETS1/pycharm_workspace/yolov5/tf2.1.1")

converter = tf.lite.TFLiteConverter.from_saved_model("/media/george/DATASETS1/pycharm_workspace/ai_learning/onnx/pb_yolov5_meshgrid_nms")

converter.target_spec.supported_ops = [

tf.lite.OpsSet.TFLITE_BUILTINS, # enable TensorFlow Lite ops.

]

converter.allow_custom_ops = True

converter.experimental_new_converter = True

# converter.optimizations = [tf.lite.Optimize.DEFAULT]

# converter.target_spec.supported_types = [tf.float16]

# converter.post_training_quantize = True

tflite_model = converter.convert()

open("y5s.tflite", "wb").write(tflite_model)

Step 7：手机端tflite模型部署。

利用Android studio 加载tensorflow官方Android detection代码

https://github.com/tensorflow/examples/tree/master/lite/examples/object_detection/android。

修改的地方也比较多，具体要看自己的模型流程是怎么样的，主要还是模型输入尺寸和输出参数是什么，修改对应的java代码即可。

如果哪里有说错或者有不清楚的可以留言.