Closed2023/07/14にクローズ2

YOLOv4TinyのONNXモデルを編集する

Python

Computer Vision

objectdetection

0y0

h5→saved model→onnxという経緯で変換したモデルについて、余分なノードを削除する

対象

以下リポジトリの環境で作成されるモデルについて取り扱う

余分なノードの削除

h5

kerasで作成された変換前のモデル
Conv2d_21を出力する過程で、LeakyRelu→Upsampleの部分がある

h5→saved model→onnx

ONNXに変換されたモデル
LeakyRelu→Upsampleまでに余分なノードが追加されてしまう
また、バージョン10以降ではUpsampleが非推奨となっているので、Resizeに置き換えたい

ONNXモデルの編集

環境

以下記事で紹介されているDocker環境を利用する
https://cyberagent.ai/blog/tech/17300/

コード

ノードの削除と置き換えを行う

import onnx
import onnx_graphsurgeon as gs
import numpy as np

def create_constants(scale_values, scale_name, roi_name):
    """スケールと空の ROI の定数テンソルを作成

    Args:
        scale_values (np.ndarray): 使用されるスケール値を含む配列
        scale_name (str): スケール テンソルの名前
        roi_name (str): ROI テンソルの名前

    Returns:
        scale_tensor: スケールの定数テンソル
        empty_roi: 空の ROI の定数テンソル
    """
    scale_tensor = gs.Constant(name=scale_name, values=scale_values)
    empty_roi = gs.Constant(name=roi_name, values=np.array([], dtype=np.float32).reshape(0))

    return scale_tensor, empty_roi

def update_model(graph):
    """特定の変更に基づいてONNXモデルを更新。

    Args:
        graph: 更新するONNXグラフ

    Returns:
        graph: 更新されたONNXグラフ
    """
    # 特定のノードを削除
    remove_op_types = ["Shape", "Gather", "Cast", "Slice", "Mul", "Div"]
    remove_node_names = ["Concat__168", "Upsample__169"]

    nodes_to_remove = [node for node in graph.nodes if node.op in remove_op_types or node.name in remove_node_names]
    graph.nodes = [node for node in graph.nodes if node not in nodes_to_remove]

    # テンソルを名前で参照する辞書を作成
    tensor_dict = {}
    for node in graph.nodes:
        for tensor in node.inputs + node.outputs:
            tensor_dict[tensor.name] = tensor

    for tensor in graph.inputs + graph.outputs:
        tensor_dict[tensor.name] = tensor

    # 新しいResizeノードを作成
    scale_tensor_1, empty_roi_1 = create_constants(np.array([1.0, 1.0, 2.0, 2.0], dtype=np.float32), "scales_tensor_1", "empty_roi_1")
    new_resize_node = gs.Node(
        op="Resize",
        inputs=[
            tensor_dict["StatefulPartitionedCall/model_1/leaky_re_lu_18/LeakyRelu:0"],
            empty_roi_1,
            scale_tensor_1
        ],
        outputs=[tensor_dict["Upsample__169:0"]],
        attrs={
            "coordinate_transformation_mode": "pytorch_half_pixel",
            "mode": "nearest"
        }
    )

    # 新しいResizeノードをグラフに追加
    graph.nodes.append(new_resize_node)

    return graph

def main():
    # 既存のモデルをロード
    model = onnx.load("yolov4_tiny.onnx")

    # オペレーターセットのバージョンを更新
    model.opset_import[0].version = 14

    # ONNXモデルを ONNX-GraphSurgeonグラフに変換
    graph = gs.import_onnx(model)

    # モデルを更新
    graph = update_model(graph)

    # 新しいONNXモデルをエクスポート
    new_model = gs.export_onnx(graph)

    # 新しいONNXモデルを保存
    onnx.save(new_model, "yolov4_tiny_modified.onnx")

if __name__ == "__main__":
    main()

ONNX（編集後）

編集後のモデル
余分なノードが削除されUpsampleがResizeに置き換わっている

テスト

ONNXモデルが正しく動作するか以下のコマンドで確認する

sit4onnx -if yolov4_tiny_modified.onnx -oep cpu

0y0

Unity Barracudaで使用する場合

Barracudaではopset=9が推奨されており、Resizeノードが使えない。
そのため、Upsampleノードは残したままプロパティを書き換える必要がある。
ただし、Barracuda内ではONNXモデルがBrracuda形式？に変換される。
（未編集のONNXモデルと編集したONNXモデルでは、変換後の形状は同じだった）
したがって、Resizeノードを使用していない場合は以下の処理は不要と思われる。
※それよりもSplitノードにAttributeを追加する処理の方が重要

import onnx
import onnx_graphsurgeon as gs
import numpy as np

def create_constants(scale_values, scale_name):
    """スケールの定数テンソルを作成

    Args:
        scale_values (np.ndarray): 使用されるスケール値を含む配列
        scale_name (str): スケール テンソルの名前

    Returns:
        scale_tensor: スケールの定数テンソル
    """
    scale_tensor = gs.Constant(name=scale_name, values=scale_values)
    return scale_tensor

def update_model(graph):
    """特定の変更に基づいてONNXモデルを更新。

    Args:
        graph: 更新するONNXグラフ

    Returns:
        graph: 更新されたONNXグラフ
    """
    # 特定のノードを削除
    remove_op_types = ["Shape", "Gather", "Cast", "Slice", "Mul", "Div"]
    remove_node_names = ["Concat__168"]

    nodes_to_remove = [node for node in graph.nodes if node.op in remove_op_types or node.name in remove_node_names]
    graph.nodes = [node for node in graph.nodes if node not in nodes_to_remove]

    # テンソルを名前で参照する辞書を作成
    tensor_dict = {}
    for node in graph.nodes:
        for tensor in node.inputs + node.outputs:
            tensor_dict[tensor.name] = tensor

    for tensor in graph.inputs + graph.outputs:
        tensor_dict[tensor.name] = tensor

    # Upsample__169のプロパティを更新
    upsample_node = next(node for node in graph.nodes if node.name == "Upsample__169")
    
    scale_tensor = create_constants(np.array([1.0, 1.0, 2.0, 2.0], dtype=np.float32), "scales_tensor")
    
    upsample_node.inputs = [tensor_dict["StatefulPartitionedCall/model_1/leaky_re_lu_18/LeakyRelu:0"], scale_tensor]
    upsample_node.attrs = {"mode": "nearest"}

    # スケール定数をグラフに追加
    graph.inputs.append(scale_tensor)

    return graph

def main():
    # 既存のモデルをロード
    model = onnx.load("yolov4_tiny.onnx")

    # オペレーターセットのバージョンを更新
    model.opset_import[0].version = 9

    # ONNXモデルを ONNX-GraphSurgeonグラフに変換
    graph = gs.import_onnx(model)

    # モデルを更新
    graph = update_model(graph)

    # 新しいONNXモデルをエクスポート
    new_model = gs.export_onnx(graph)

    # 新しいONNXモデルを保存
    onnx.save(new_model, "yolov4_tiny_modified.onnx")

if __name__ == "__main__":
    main()

このスクラップは2023/07/14にクローズされました