[英]How to implement custom output layer with dynamic shape in Keras?
我想用 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 加載了預訓練的權重並運行了model.predict
,但是 output 給了我一個錯誤:
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 的情況下運行了這個 model,並使用相同的 function 但單獨修改了它的輸出,它工作正常,但它不使用占位符。 我應該怎么做才能實現這一目標?
好的,我自己發現的。 所有必須做的是:
outputs = outputs.write(out_idx, image_out)
聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.