(beta) onnxの標準機能で２つ以上のonnxファイルを１つにマージする

docker run --gpus all -it --rm \
-v `pwd`:/home/user/workdir \
ghcr.io/pinto0309/openvino2tensorflow:latest

git clone https://github.com/protocolbuffers/protobuf.git
cd protobuf
git checkout v3.16.0
git submodule update --init --recursive
mkdir build_source && cd build_source
cmake ../cmake \
-Dprotobuf_BUILD_SHARED_LIBS=OFF \
-DCMAKE_INSTALL_PREFIX=/usr \
-DCMAKE_INSTALL_SYSCONFDIR=/etc \
-DCMAKE_POSITION_INDEPENDENT_CODE=ON \
-Dprotobuf_BUILD_TESTS=OFF \
-DCMAKE_BUILD_TYPE=Release
make -j$(nproc)
sudo make install

cd ../..

export CMAKE_ARGS="-DONNX_USE_PROTOBUF_SHARED_LIBS=ON"
git clone --recursive https://github.com/onnx/onnx.git
cd onnx
git checkout 94fb49259018fc9ed95732bd7eb609bd518b0ee8
sudo pip3 install -e .

PINTO

hep_lum_lol_HxW.onnx
    input names: input.1
    output names: 227, 233

hep_ndm_lol_encoder_HxW.onnx
    input names: input.1
    output names: 901

compose_onnx_file.py

import onnx

model1 = onnx.load('hep_lum_lol_HxW.onnx')
model1 = onnx.compose.add_prefix(model1, prefix='lum_')

model2 = onnx.load('hep_ndm_lol_encoder_HxW.onnx')
model2 = onnx.compose.add_prefix(model2, prefix='ndm_enc_')
combined_model1 = onnx.compose.merge_models(
    model1, model2,
    io_map=[('lum_227', 'ndm_enc_input.1')]
)

onnx.save(combined_model1, 'combined.onnx')

PINTO

onnx３連マージのサンプル

hep_lum_lol_HxW.onnx
    input names: input.1
    output names: 227, 233

hep_ndm_lol_encoder_HxW.onnx
    input names: input.1
    output names: 901

hep_ndm_lol_decoder_HxW.onnx
    input names: input.1
    output names: 901

onnx_triple_merge.py

import onnx

model1 = onnx.load('hep_lum_lol_HxW.onnx')
model1 = onnx.compose.add_prefix(model1, prefix='lum_')

model2 = onnx.load('hep_ndm_lol_encoder_HxW.onnx')
model2 = onnx.compose.add_prefix(model2, prefix='ndm_enc_')
combined_model1 = onnx.compose.merge_models(
    model1, model2,
    io_map=[('lum_227', 'ndm_enc_input.1')]
)

model3 = onnx.load('hep_ndm_lol_decoder_HxW.onnx')
model3 = onnx.compose.add_prefix(model3, prefix='ndm_dec_')
combined_model2 = onnx.compose.merge_models(
    combined_model1, model3,
    io_map=[('ndm_enc_901', 'ndm_dec_input.1')]
)

onnx.save(combined_model2, 'hep_combine_lol_HxW.onnx')

PINTO