Adding argmax causes a prediction error

[geoexploring]

Description

The problem is exactly the same as #1025. Adding the argmax layer to the onnx model will cause the prediction to fail, but without adding argmax, running argmax on the cpu will be very slow.

Environment

TensorRT Version: 8.0.1.6
NVIDIA GPU: RTX2080
NVIDIA Driver Version: 11.6
CUDA Version: 10.2
CUDNN Version: 7.6
Operating System: WIN10

Thanks.

#2416 (comment)

[zerollzeng]

Can you try the latest release of TensorRT? like 8.4 or 8.5? If there are still some problems can you provide a reproduce to us?
Thanks!

[geoexploring]

Thank you for your advice. When I changed TensorRT-8.0.1.6 to TensorRT-8.4.2.4, an error was reported:

[2022-10-26 10:01:43][error][trt_builder.cpp:34]:NVInfer: 1: [stdArchiveReader.cpp::nvinfer1::rt::StdArchiveReader::StdArchiveReader::35] Error Code 1: Serialization (Serialization assertion safeVersionRead == safeSerializationVersion failed.Version tag does not match. Note: Current Version: 0, Serialized Engine Version: 43)
[2022-10-26 10:01:43][error][trt_builder.cpp:34]:NVInfer: 4: [runtime.cpp::nvinfer1::Runtime::deserializeCudaEngine::50] Error Code 4: Internal Error (Engine deserialization failed.)

Similar to #1736.

[zerollzeng]

Please rebuild the engine with TRT 8.4

[geoexploring]

Thanks! Internal Error was fixed.

But the problem of prediction failure after adding argmax remains. I've been pretty busy lately. I'll provide a reproduction to you when I'm free.

[geoexploring]

@zerollzeng , for the solution, please see Wrong result in TensorRT, but it seems something is working correctly.