softmax + argmax takes too much time

[Xiao-Hu-Z]

I would like to ask, considering using TensorRT plugin or CUDA Kernel to implement more efficient Argmax operations, will it be faster?

[lix19937]

Yes, use TensorRT plugin will be more efficient.

[geraldstanje]

hi @lix19937 or @Xiao-Hu-Z - im not sure if my question relates to your question.
i also have a post-processing step in tensorRT (using python) to convert probabilities into a class (for binary classification) and than to a label. can that be done with TensorRT plugin too? I'm looking for ways to accelerate it.
also, is i have a pre-processing step (tokenizer in python) - is there a way to accelerate that too using c++ or rust for example?

[yuanyao-nv]

cc @samurdhikaru @venkywonka for thoughts on plugin perf