TensorFlow 不需要的子图

Question

我正在尝试构建一个 TFLite 模型，该模型由用于检测的高效 det 和用于识别的 keras-ocr（只是识别器模型）组成，我需要它在移动设备上运行得快一些。 我想完全量化模型，但在Netron 中检查模型时发现模型有一些子图，而 TFLite 不支持使用子图对模型进行量化。

List of subgraphs in model

这是我正在使用的功能模型定义以及我认为生成子图的层。

MAX_DETECTIONS = 10

img_inputs = tf.keras.Input(shape=(1080, 1920, 3), batch_size = 1)
boxes, scores, classes, _ = efficient_det_model_keras.model(img_inputs)
crops = CroppingLayer(MAX_DETECTIONS)(img_inputs, boxes, scores)
recognitions = RecognitionLayer(recognizer.prediction_model, MAX_DETECTIONS)(crops)

model = tf.keras.Model(inputs=img_inputs, outputs=(recognitions, scores), name="alpr")

class CroppingLayer(tf.keras.layers.Layer):
    def __init__(self, MAX_RECOGNITIONS):
        super(CroppingLayer, self).__init__()
        self.MAX_RECOGNITIONS = MAX_RECOGNITIONS
        self.threshold = 0.3
    
    def get_crop(self, input_image, ymin, xmin, height, width):
        cropped = tf.image.crop_to_bounding_box(input_image, ymin, xmin, height, width)
        cropped = tf.image.resize(cropped, (31, 200))

        cropped = tf.image.rgb_to_grayscale(cropped)/255
        
        return cropped
    
    
    def call(self, input_image, boxes, scores):
        res = []

        ymin = tf.cast(boxes[0, :, 0], tf.int32)
        xmin = tf.cast(boxes[0, :, 1], tf.int32)
        ymax = tf.cast(boxes[0, :, 2], tf.int32)
        xmax = tf.cast(boxes[0, :, 3], tf.int32)
        height = tf.math.abs(ymax - ymin)
        width = tf.math.abs(xmax - xmin)

        for i in range(self.MAX_RECOGNITIONS):
            cropped = self.get_crop(input_image[0], ymin[i], xmin[i], height[i], width[i])

            res.append(cropped)

        batch = [tf.stack(res, axis = 0)]

        return tf.stack(batch, axis = 0)

有什么办法可以摆脱这些子图来制作单个图模型？

Answer 1

看起来 for 循环语句导致在 TensorFlow 图中创建了一个控制流操作。如果识别数可以固定，解开 for 循环语句可能有助于创建单个子图。