Feature: Transpose optimization for Softmax

[fthielke]

Thanks to the optimizations by tf2onnx, most Transpose operations in converted models are eliminated.
However, Softmax does not seem to be accounted for yet.

Take the following simple SavedModel:

import tensorflow as tf
from keras import Model
from keras.layers import Input, Softmax
import keras.backend as K
K.set_image_data_format('channels_first')

inputs = Input((1, None, None), name='input')
outputs = Softmax(axis=1, name='output')(inputs)
m = Model(inputs, outputs)

tf.saved_model.save(
    obj=m,
    export_dir='test_softmax'
)

When exporting it via python -m tf2onnx.convert --saved-model test_softmax --output test_softmax/model.onnx --tag serve --signature_def serving_default --opset 14, the result looks like this:

Same for every other channels_first model I export - the Softmax operation is always surrounded by Transposes.

[TomWildenhain-Microsoft]

@fthielke Yes, this would be a good optimization to have. Do you think you could try to add it to the transpose optimizer and submit a PR?

tensorflow-onnx/tf2onnx/optimizer/transpose_optimizer.py

Line 224 in 446494e

"Sigmoid": self._simple_through_handler,

[fatcat-z]

Wait for user to fix related PR issue.