如何在OpenVINO模型优化器中添加FusedBatchNormV3支持？

Let's break down how to resolve the FusedBatchNormV3 conversion error and implement the ScaleShift-based replacement you're targeting.

First, Understand the Core Issue #

The FusedBatchNormV3 op from your TensorFlow 1.5/1.3 ResNet50 isn't natively supported by the version of OpenVINO Model Optimizer you're using. Since you want to map it to ScaleShift (which is supported and aligns with the normalized batch norm behavior for frozen models), we'll use Model Optimizer's custom subgraph replacement feature to handle this.

Step 1: Locate the ScaleShift Operation Class #

In OpenVINO's Model Optimizer, the ScaleShift op is implemented in the extensions library. You can import it directly in your custom replacement script from:

from extensions.ops.scale_shift import ScaleShift

This class handles the scale (multiplication) and shift (addition) operations that match the frozen batch normalization behavior (since mean/variance are constants after freezing).

Step 2: Create a Custom Replacement Script #

Create a Python script (e.g., fused_bn_v3_to_scaleshift.py) with the following logic to replace FusedBatchNormV3 nodes:

import numpy as np
from mo.front.common.partial_infer.utils import int64_array
from mo.front.tf.replacement import FrontReplacementSubgraph
from mo.graph.graph import Graph
from extensions.ops.scale_shift import ScaleShift

class FusedBatchNormV3ToScaleShift(FrontReplacementSubgraph):
    enabled = True

    def pattern(self):
        # Match the FusedBatchNormV3 node structure in frozen TensorFlow graphs
        return dict(
            nodes=[
                ('fused_bn', dict(op='FusedBatchNormV3')),
                ('gamma', dict(op='Const')),
                ('beta', dict(op='Const')),
                ('mean', dict(op='Const')),
                ('variance', dict(op='Const')),
            ],
            edges=[
                ('gamma', 'fused_bn', {'in': 1}),
                ('beta', 'fused_bn', {'in': 2}),
                ('mean', 'fused_bn', {'in': 3}),
                ('variance', 'fused_bn', {'in': 4}),
            ]
        )

    def replace_sub_graph(self, graph: Graph, match: dict):
        fused_bn_node = match['fused_bn']
        gamma_node = match['gamma']
        beta_node = match['beta']
        mean_node = match['mean']
        variance_node = match['variance']

        # Extract constant values from the nodes
        gamma = gamma_node.value
        beta = beta_node.value
        mean = mean_node.value
        variance = variance_node.value
        epsilon = fused_bn_node.epsilon if hasattr(fused_bn_node, 'epsilon') else 1e-5

        # Calculate ScaleShift parameters: scale = gamma / sqrt(variance + epsilon)
        # shift = beta - (gamma * mean) / sqrt(variance + epsilon)
        std = np.sqrt(variance + epsilon)
        scale = gamma / std
        shift = beta - (gamma * mean) / std

        # Create new ScaleShift node
        scaleshift_node = ScaleShift(graph, {
            'name': fused_bn_node.name + '_replaced',
            'axis': 3,  # For NHWC format (TensorFlow default)
            'output_shape': fused_bn_node.out_node().shape,
        }).create_node()

        # Reconnect edges: input to ScaleShift, and redirect output
        fused_bn_node.in_node(0).out_port(0).connect(scaleshift_node.in_port(0))
        fused_bn_node.out_port(0).get_connection().set_source(scaleshift_node.out_port(0))

        # Add scale and shift as constant inputs to ScaleShift
        graph.add_node(
            name=fused_bn_node.name + '_scale',
            op='Const',
            value=scale,
            shape=int64_array(scale.shape),
            kind='op'
        ).out_port(0).connect(scaleshift_node.in_port(1))

        graph.add_node(
            name=fused_bn_node.name + '_shift',
            op='Const',
            value=shift,
            shape=int64_array(shift.shape),
            kind='op'
        ).out_port(0).connect(scaleshift_node.in_port(2))

        # Remove old nodes (optional, cleanup)
        graph.remove_nodes_from([fused_bn_node.id, gamma_node.id, beta_node.id, mean_node.id, variance_node.id])

Step 3: Run Model Optimizer with the Custom Extension #

Modify your MO command to include the custom extension script:

sudo python3 mo.py \
 --input_model ~<PATH_TO_MODEL>/model.pb \
 --output_dir ~<PATH_TO_MODEL> \
 --data_type FP16 -b 1 \
 --extensions ./fused_bn_v3_to_scaleshift.py

Step 4: Fix the Float16 Clipping Error #

The initial infinity clipping error happens because some values exceed FP16's range during conversion. To mitigate this:

1. First run the conversion with --data_type FP32 to verify the replacement works without overflow.
2. If FP32 works, adjust the scale/shift calculation in your replacement script to clamp values to FP16's valid range (e.g., scale = np.clip(scale, -65504, 65504) for FP16).

Additional Notes #

• If you can upgrade TensorFlow to 1.15 (the last TF 1.x release), the exported ResNet50 might use FusedBatchNormV2 which is natively supported by newer OpenVINO versions, skipping the need for custom replacement.
• Ensure your OpenVINO version is compatible with Ubuntu 18.04 (note: you wrote 18.03, which is a typo—Ubuntu 18.04 is the LTS release).

内容的提问来源于stack exchange，提问作者tuscasp