如何在Keras中实现具有动态形状的自定义输出层？

Question

我想用Tensorflow 2.0后端在Keras中实现YOLO tiny。我想创建一个新的自定义YoloLayer，它对前一层的输出执行非最大值抑制，并生成具有(batch_size, num, 6)形状的张量，其中num是许多已找到的预测，每个预测都表示为[x, y, w, h, prob, class]。我还在__init__()方法中设置了self.trainable = False。以下是我的call方法：

def call(self, inputs, **kwargs):
        predictions = inputs[...,:5]
        x = tf.math.add(self.cols, tf.nn.sigmoid(predictions[...,0])) / self.grid_size # x
        y = tf.math.add(self.rows, tf.nn.sigmoid(predictions[...,1])) / self.grid_size # y
        w = tf.multiply(self.anchors_w, tf.math.exp(predictions[...,2])) / self.grid_size # w
        h = tf.multiply(self.anchors_h, tf.math.exp(predictions[...,3])) / self.grid_size # h
        c = tf.nn.sigmoid(predictions[...,4]) # confidence
        
        bounds = tf.stack([x, y, w, h], -1)
        classes = inputs[...,5:]
        probs = tf.multiply(tf.nn.softmax(classes), tf.expand_dims(c, axis=-1))
        prob_mask = tf.greater(probs, self.threshold)
        suppressed_indices = tf.where(prob_mask)
        suppressed_probs = tf.gather_nd(probs, suppressed_indices[...,:3])
        suppressed_boxes = tf.gather_nd(bounds, suppressed_indices[...,:3])
        box_coords = tf.stack([
            suppressed_boxes[...,1] - suppressed_boxes[...,3] / 2., #y1
            suppressed_boxes[...,0] - suppressed_boxes[...,2] / 2., #x1
            suppressed_boxes[...,1] + suppressed_boxes[...,3] / 2., #y2
            suppressed_boxes[...,0] + suppressed_boxes[...,2] / 2., #x2
        ], axis=-1)

        out = tf.TensorArray(tf.float32, size=0, dynamic_size=True)

        for i in range(tf.shape(inputs)[0]):
            image_out = tf.TensorArray(tf.float32, size=self.classes)
            for c in range(self.classes):
                class_probs = suppressed_probs[i,:,c]
                indices = tf.image.non_max_suppression(box_coords[i], class_probs, 10,
                                                       iou_threshold=self.nms_threshold,
                                                       score_threshold=self.threshold)
                
                if tf.size(indices) > 0:
                    final_probs = tf.expand_dims(tf.gather(class_probs, indices), axis=-1)
                    final_boxes = tf.gather(suppressed_boxes[i], indices)
                    class_vec = tf.ones((tf.shape(final_probs)[0], 1)) * c
                    image_out.write(c, tf.concat([final_boxes, final_probs, class_vec], axis=1))
                    
            image_out = image_out.concat()
            out.write(i, image_out)
        
        out = out.stack()
        return out

然后，model.summary()返回：

Model: "sequential_1"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
...
_________________________________________________________________
yolo_layer (YoloLayer)       (None, None, 6)           0         
=================================================================
...

我加载了此模型的预先训练权重并运行model.predict，但输出给我一个错误：

InvalidArgumentError:  Tried to stack elements of an empty list with non-fully-defined element_shape: [?,6]
     [[node sequential_1/yolo_layer/TensorArrayV2Stack/TensorListStack (defined at <ipython-input-2-fbae137dd1a2>:96) ]] [Op:__inference_predict_function_4604]

我也在没有YoloLayer的情况下运行了这个模型，并使用相同的函数修改了它的输出，但是单独的，它工作正常，但没有占位符。我该怎么做才能做到这一点

Answer 1

好的，我自己发现的。必须做的是：

outputs = outputs.write(out_idx, image_out)