YOLOv5 Android(完结)_yolov5s_relu.onnx

最近在搞移动端的模型部署包括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代码即可。

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