TensorRT へ ArgMax/ArgMin を適用する場合に INT64 や INT32 が入力テンソルに指定されていると変換エラーとなる件へのワークアラウンド

https://github.com/PINTO0309/PINTO_model_zoo/issues/246

1. ArgMax の入力テンソルのタイプを INT64 のままでONNXへエクスポートした場合

   • 生成されたONNXを onnx-simplifier で最適化しようとすると [ONNXRuntimeError] : 9 : NOT_IMPLEMENTED : Could not find an implementation for ArgMax(11) node with name 'ArgMax_xxxx' が発生して Abort する

2. INT64 から INT32 へ Cast したうえでONNXへエクスポートした場合

   • ArgMax を含むONNXファイルの生成は正常終了する
   • 生成されたONNXの onnx-simplifier による最適化も成功する
   • onnx-simplifier 適用後に trtexec あるいは onnx2trt を使用して TensorRT Engine へ変換しようとすると下図のエラーが発生して変換に失敗する

   While parsing node number 1157 [ArgMax -> "3084"]:
   --- Begin node ---
   input: "3083"
   output: "3084"
   name: "ArgMax_2581"
   op_type: "ArgMax"
   attribute {
     name: "axis"
     i: -1
     type: INT
   }
   attribute {
     name: "keepdims"
     i: 0
     type: INT
   }
   
   --- End node ---
   ERROR: onnx2trt_utils.cpp:54 In function argMinMaxHelper:
   [8] Assertion failed: (tensor.getType() != nvinfer1::DataType::kINT32) && "This version of TensorRT does not support INT32 ArgMin/ArgMax input data."
   

3. INT64 から Float64 へ Cast したうでONNXへエクスポートした場合

   • ArgMax を含むONNXファイルの生成は正常終了する
   • 生成されたONNXの onnx-simplifier による最適化も成功する
   • onnx-simplifier で最適化した ONNX の TensorRT Engine への変換も成功する

   onnx2trt clip_vit_b_32.onnx -o clip_vit_b_32.trt -b 1 -d 16 -v
   
   [2022-04-02 13:07:15 WARNING] TensorRT was linked against cuBLAS/cuBLAS LT 11.6.5 but loaded cuBLAS/cuBLAS LT 11.6.1
   [2022-04-02 13:07:15    INFO] [MemUsageChange] Init cuBLAS/cuBLASLt: CPU +0, GPU +8, now: CPU 4041, GPU 2430 (MiB)
   [2022-04-02 13:07:15    INFO] [MemUsageChange] Init cuDNN: CPU +0, GPU +8, now: CPU 4041, GPU 2438 (MiB)
   [2022-04-02 13:07:15    INFO] [MemUsageChange] TensorRT-managed allocation in building engine: CPU +4, GPU +289, now: CPU 4, GPU 289 (MiB)
   Writing TensorRT engine to clip_vit_b_32.trt
   All done
   

Hello, I would like to know at which step you convert int to float. Is it in the PyTorch model where you set the input as float, then add a layer to convert it to int?