diff options
Diffstat (limited to 'model-optimizer/extensions/ops/priorbox.py')
| -rw-r--r-- | model-optimizer/extensions/ops/priorbox.py | 20 |
1 files changed, 18 insertions, 2 deletions
diff --git a/model-optimizer/extensions/ops/priorbox.py b/model-optimizer/extensions/ops/priorbox.py index 0caf0e957..e494097d2 100644 --- a/model-optimizer/extensions/ops/priorbox.py +++ b/model-optimizer/extensions/ops/priorbox.py @@ -52,7 +52,7 @@ class PriorBoxOp(Op): 'step', 'step_h', 'step_w', - 'offset' + 'offset', ] def backend_attrs(self): @@ -71,6 +71,22 @@ class PriorBoxOp(Op): def priorbox_infer(node: Node): layout = node.graph.graph['layout'] data_shape = node.in_node(0).shape - num_ratios = ((node.flip + 1) * len(node.aspect_ratio) + 1) * len(node.min_size) + len(node.max_size) + + # calculate all different aspect_ratios (the first one is always 1) + # in aspect_ratio 1/x values will be added for all except 1 if flip is True + ar_seen = [1.0] + ar_seen.extend(node.aspect_ratio.copy()) + if node.flip: + for s in node.aspect_ratio: + ar_seen.append(1.0/s) + + ar_seen = np.unique(np.array(ar_seen).round(decimals=6)) + + num_ratios = 0 + if len(node.min_size) > 0: + num_ratios = len(ar_seen) * len(node.min_size) + + num_ratios = num_ratios + len(node.max_size) + res_prod = data_shape[get_height_dim(layout, 4)] * data_shape[get_width_dim(layout, 4)] * num_ratios * 4 node.out_node(0).shape = np.array([1, 2, res_prod], dtype=np.int64) |