转yolov5模型时，不转后处理部分

最新推荐文章于 2024-06-14 19:05:55 发布

原创

最新推荐文章于 2024-06-14 19:05:55 发布 · 693 阅读

5 ·

CC 4.0 BY-SA版权

文章标签：

#YOLO #转换模型

本文讲述了如何在将后处理部分从硬件转移到CPU以支持不被硬件直接处理的情况下的模型转换，涉及TFLite和ONNX版本的TensorFlow模型中Detect类的相应代码修改，以及坐标变换的处理策略。

一、背景

由于部署在硬件上的时候，后处理部分硬件处理不支持，需要挪到cpu上处理。

二、转int8.tflite版本时挪出后处理部分

需要修改的文件models/tf.py
将TFDetect(keras.layers.Layer)的call()函数修改为下面部分

    def call(self, inputs):
        print('************* deploy ******************')
        z = []  # inference output
        x = []
        for i in range(self.nl):
        	###  原始 ####
            # x.append(self.m[i](inputs[i])) 
            ######## 新增 ########
            if True: 
                temp = self.m[i](inputs[i])
                z.append(tf.reshape(temp,[-1,6]))
                # print('shape', self.m[i](inputs[i]).reshape(-1,6).shape)
            continue
            ######## 新增 ########
            # x(bs,20,20,255) to x(bs,3,20,20,85)
            ny, nx = self.imgsz[0] // self.stride[i], self.imgsz[1] // self.stride[i]
            x[i] = tf.reshape(x[i], [-1, ny * nx, self.na, self.no])

            if not self.training:  # inference
                y = x[i]
                grid = tf.transpose(self.grid[i], [0,