ERROR: UffParser: Parser error: softmax/Max: Invalid reduction axes

I trained an OCR model using Keras and then I converted to .pb file and this model has a mean layer that I defined this layer in Keras model as follows:

...
x = BatchNormalization()(x)
x = Lambda(lambda x: mean(x, axis=2))(x)
...

I want to convert .pb model to pure TensorRT, but in the UFF parser I got this error:

[TensorRT] INFO: UFFParser: parsing lambda_1/Mean
[TensorRT] INFO: UFFParser: Applying order forwarding to: lambda_1/Mean
[TensorRT] INFO: UFFParser: parsing softmax/Max
[TensorRT] ERROR: UffParser: Parser error: softmax/Max: Invalid reduction axes

My input shape in the TensorFlow and Keras model is : (256,128,3), but when I define this shape here, I got the error, but when I define (3,256,128), I don’t get the error, then I have to define this shape for the parser.

I see the operation Mean is defined the Reduce op in UFF.

I test on jetson nano:
Tensorflow : 1.12.2
jetpack : 4.2.2
TensorRT 5.x

I did these follow steps:

dynamic_graph = gs.DynamicGraph('ocr.ob')
_ = uff.from_tensorflow(
    dynamic_graph.as_graph_def(),
    output_nodes=['softmax/truediv'],
    output_filename='my_ocr.uff',
    text=True,
    debug_mode=DEBUG_UFF)

The .uff and .pbtxt files created correctly.

with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.UffParser() as parser:
    builder.max_workspace_size = 1 << 28
    builder.max_batch_size = 1
    builder.fp16_mode = True

    # parser.register_input('Input', INPUT_DIMS)
    # parser.register_output('MarkOutput_0')
    # parser.parse(spec['tmp_uff'], network)

    parser.register_input('the_input', (3,256,128))
    parser.register_output('MarkOutput_0')
    parser.parse('my_ocr.uff', network)

I found the softmax/Max from .pbtxt generated file, and the part of that node as follows:

nodes {
    id: "softmax/sub"
    inputs: "batch_normalization_6/FusedBatchNorm"
    inputs: "softmax/Max"
    operation: "Binary"
    fields {
      key: "func"
      value {
        s: "sub"
      }
    }
  }
  nodes {
    id: "softmax/Max"
    inputs: "batch_normalization_6/FusedBatchNorm"
    operation: "Reduce"
    fields {
      key: "axes"
      value {
        i_list {
          val: -1
        }
      }
    }
    fields {
      key: "func"
      value {
        s: "max"
      }
    }
    fields {
      key: "keepdims"
      value {
        b: true
      }
    }
  }

In My opinion, I likely use the plugin Reduce operation from UFF Operations instead of Mean node/operation, but I don’t know How do i do this?

If answer the question is hard, You suppose the trained network have some layers and I want to replace ones of nodes(like Conv) with corresponding UFF node/operation ? How do I do?

Hi,

Please refer to the below samples:
https://github.com/NVIDIA/TensorRT/tree/release/7.0/samples/opensource/sampleUffSSD
https://github.com/NVIDIA/TensorRT/tree/release/7.0/samples/opensource/sampleUffPluginV2Ext

Thanks