Zenn
Open5

MAXIM の最適化チャレンジ

TensorFlow
#
onnx
PINTOPINTO
python -m tf2onnx.convert \
--opset 13 \
--inputs input_1:0[1,256,256,3] \
--inputs-as-nchw input_1:0 \
--saved-model . \
--output maxim_s-2_dehazing_sots-outdoor_1.onnx
PINTOPINTO

特殊OP ScaleAndTranslate が含まれていて onnxsim が Abort するので自力で形状推定を試みる。まずは、形状推定が失敗して onnx2tf が Abort する箇所を重点的に対処する。モデルの序盤で shape を使用した形状推定があるので静的な形状を特定できていない。したがって、ScaleAndTranslate が登場する前までの最小限のモデルへ分離して最適化を施してみて、正しい形状を知ることから始める。

# 形状推定が失敗する最初の Mul で分離する
sne4onnx \
--input_onnx_file_path maxim_s-2_dehazing_sots-outdoor_1.onnx \
--input_op_names input_1:0 \
--output_op_names StatefulPartitionedCall/s2_model/tf.math.multiply/Mul:0 StatefulPartitionedCall/s2_model/tf.math.multiply_1/Mul:0 \
--output_onnx_file_path maxim_s-2_dehazing_sots-outdoor_1_cut1.onnx

# 出力形状を調べる
sit4onnx -if maxim_s-2_dehazing_sots-outdoor_1_cut1.onnx -oep cpu

INFO: file: maxim_s-2_dehazing_sots-outdoor_1_cut1.onnx
INFO: providers: ['CPUExecutionProvider']
INFO: input_name.1: input_1:0 shape: [1, 3, 256, 256] dtype: float32
INFO: test_loop_count: 10
INFO: total elapsed time:  755.549430847168 ms
INFO: avg elapsed time per pred:  75.5549430847168 ms
INFO: output_name.1: StatefulPartitionedCall/s2_model/tf.math.multiply_1/Mul:0 shape: [1, 256, 256, 32] dtype: float32
INFO: output_name.2: StatefulPartitionedCall/s2_model/tf.math.multiply/Mul:0 shape: [1, 256, 256, 32] dtype: float32
PINTOPINTO

次に ConvTranspose の output_shape が None なおかつ 入力形状不定、出力形状不定のオペレーションで形状推定に失敗するのでモデルを分断してテスト推論をした結果を output_shape に設定する。

sne4onnx \
--input_onnx_file_path maxim_s-2_dehazing_sots-outdoor_1.onnx \
--input_op_names input_1:0 \
--output_op_names StatefulPartitionedCall/s2_model/stage_0_decoder_block_2_ConvTranspose_0/BiasAdd:0 \
--output_onnx_file_path maxim_s-2_dehazing_sots-outdoor_1_cut2.onnx

sit4onnx -if maxim_s-2_dehazing_sots-outdoor_1_cut2.onnx -oep cpu

INFO: file: maxim_s-2_dehazing_sots-outdoor_1_cut2.onnx
INFO: providers: ['CPUExecutionProvider']
INFO: input_name.1: input_1:0 shape: [1, 3, 256, 256] dtype: float32
INFO: test_loop_count: 10
INFO: total elapsed time:  2592.797040939331 ms
INFO: avg elapsed time per pred:  259.2797040939331 ms
INFO: output_name.1: StatefulPartitionedCall/s2_model/stage_0_decoder_block_2_ConvTranspose_0/BiasAdd:0 shape: [1, 128, 64, 64] dtype: float32
PINTOPINTO
sne4onnx \
--input_onnx_file_path maxim_s-2_dehazing_sots-outdoor_1.onnx \
--input_op_names input_1:0 \
--output_op_names StatefulPartitionedCall/s2_model/stage_0_cross_gating_block_2_ConvTranspose_0/BiasAdd:0 \
--output_onnx_file_path maxim_s-2_dehazing_sots-outdoor_1_cut3.onnx

sit4onnx -if maxim_s-2_dehazing_sots-outdoor_1_cut3.onnx -oep cpu

INFO: file: maxim_s-2_dehazing_sots-outdoor_1_cut3.onnx
INFO: providers: ['CPUExecutionProvider']
INFO: input_name.1: input_1:0 shape: [1, 3, 256, 256] dtype: float32
INFO: test_loop_count: 10
INFO: total elapsed time:  3196.537733078003 ms
INFO: avg elapsed time per pred:  319.6537733078003 ms
INFO: output_name.1: StatefulPartitionedCall/s2_model/stage_0_cross_gating_block_2_ConvTranspose_0/BiasAdd:0 shape: [1, 128, 64, 64] dtype: float32