TensorRT dimension order rearrangement after operation

Autonomous Machines Jetson & Embedded Systems Jetson AGX Xavier
1

I am not sure whether this behavior is a bug or not.

The main architecture of my model is ResNet50. I try to connect the output tensor of conv5_block3_out (I use the name in Keras shown in model.summary) to another computation subgraph. I observe some weird phenomenon in the subgraph. I name the tensor output from last convolution block “feature_map”. According to ResNet50 architecture and NCHW order in the TensorRT model, the dimension should be (2048, 15, 20).

Here is the brief introduction of my subgraph and use pythonic TensorFlow code for assistance.

  • I connect the "feature_map" to a "reduce_mean" node.
reduce_mean = tf.reduce_mean(feature_map, [1, 2])
  • "reduce_mean" then flow in my customized plugin.
    • output_tensor = my_plugin(reduce_mean)
  • I substact the "ouptut_tensor" to "feature_map"
    • final_tensor = output_tensor - feature_map

    The customized plugin is written in C++ which I specify to output dims would be (2048, 15, 20). After I logout the input dims in the customized plugin from reduce_mean, it shows dimensions are (1, 1, 2048) instead of (2048, 1, 1). This is the first time I observe the dimension order rearrangement. Then errors prompt that the dimension is mismatched in subtraction, which is (2048, 15, 20) and (15, 20, 2048). I believe (15, 20, 2048) is the dimension from “feature_map”. Here comes the even stranger thing. How come an operand’s dimension would change after those two operations? I reasonably doubt the dimensions would automatically be sorted after operations. I consider this is not the expected behavior especially in TensorRT which dimensions really matters.

    Btw should the dims include a batch dim?

    2

    Hi,

    I’m afraid that your use case may not be supported.
    Do you meet any error when converting uff model into TRT engine?

    I tried your use case and it doesn’t support.

    [TensorRT] ERROR: UFFParser: Parser error: Mean: Invalid reduction axes

    Thanks.

    3

    Actually, I didn’t meet any issue while conversion. I suppose the reason of inconsistency is because these nodes are just part of the subgraph and TensorRT automatically select different nodes to fuse. Sorry for not mentioning the simplification. I do it for making the description as easy as possible. Is there any method to assure we get the same output from TensorRT? Moreover, the “invalid reduction axes” in the error message doesn’t make sense. Same axes arguments work pretty fine in Tensorflow code.

    4

    Hi,

    We try to reproduce your issue but get stuck in the “Invalid reduction axes” error.
    Would you mind to share your source for us verifying?

    Thanks.