argmax bug with tensorrt 5.1


I tried to use a uff model converted by tensorrt 5.1 from tensorflow, but I found a bug (might be a bug) with argmax. In tensorflow, the axis of doing argmax could be set, and it is usually the last dimension of the input tensor (axis = 3). Since the order of tensorflow model is NHWC, whereas, the order of tensorrt is NCHW, the axis points out to the last dimension (axis = 3) in tensorflow should be changed to be second dimension in tensorrt (axis = 2), however, when I tried to print out the extent of each dimension of the output of argmax in tensorrt, it returned the dimensions by channel and height, but it should return height and width.

It seems like the axis to be performed argmax in tensorflow is copied directly to tensorrt, and there is no axis change (axis = 3 in tensorflow should be changed to axis = 1 in tensorrt, considering the first axis = 0) during the model conversion from .pb to .uff.

In other words, the argmax should be performed in the axis related to channel and not in the axis related to width.

I think this only happens when argmax performs with the axis is not 0, and it due to the order change from NHWC to NCHW.

If I made any mistakes, please advise me.

Thank you.

I found the same issue, did you find a solution?