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

Question

我想在Tensorflow 2.0后端在Keras中实现YOLO-tiny。我想制作一个新的自定义YoloLayer，它对上一层的输出执行非最大抑制，并使张量的形状为(batch_size, num, 6)，其中num是找到的多个预测，每个预测表示为{ {1}}。我还在[x, y, w, h, prob, class]方法中设置了self.trainable = False。这是我的__init__()方法：

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

我还运行了没有YoloLayer的该模型，并使用相同的功能但分开地修改了其输出，并且它的工作原理正确，但不包含占位符。我该怎么做才能做到这一点？

Answer 1

好的，我自己发现了。所有要做的是：

outputs = outputs.write(out_idx, image_out)