rembg使用onnxruntime-gpu总结

针对不同显卡在cuda11.x和12.x上使用的分享

安装rembg[gpu,cli]时
1
环境：

       p40 24G，
       Ubuntu 22.04.4 LTS
       python 3.10.12
       NVIDIA-SMI： 535.183.06，
       nvcc： release 12.2, V12.2.140，
       cudnn：9.2.1
       部署后，需要部署TensorRT，部署TensorRT-10.2.0.19
       onnxruntime-gpu：1.18.1，安装的适配cuda12.x的版本

运行：rembg i -m u2net dog.jpg dog-u2net.jpg
结果：TensorRT报错
onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 : FAIL : TensorRT EP failed to create engine from network for fused node: TensorrtExecutionProvider_TRTKernel_graph_torch-jit-export_9754103867505355785_0_0
预估：可能与onnxruntime安装的支持cuda12.x的版本有关
解决： 重新安装rembg:

   pip install onnxruntime-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/
   pip install rembg[gpu] --no-deps
   pip install numpy opencv-python-headless pillow pooch pymatting scikit-image scipy tqdm

安装后确认：

   import  onnxruntime as ort
   print(f"onnxruntime device: {ort.get_device()}") # output: GPU
   print(f'ort avail providers: {ort.get_available_providers()}') # output: ['CUDAExecutionProvider', 'CPUExecutionProvider']
   ort_session = ort.InferenceSession('/root/.u2net/u2net.onnx', providers=["CUDAExecutionProvider"])
   print(ort_session.get_providers())

通过代码运行：

   python3 -c "from rembg import remove, new_session; from PIL import Image; output = remove(Image.open('dog.jpg'), session=new_session('u2net', ['CUDAExecutionProvider'])); output.save('dog-u2net.jpg')"

通过代码运行后正常，运行中gpu被使用，但是通过命令行运行依然报上面的错误，主要还是在代码中集成居多，命令行问题预估是版本不一致可能性很高。

2
环境：

       a100 40G，
       python 3.10.12
       Ubuntu 22.04.3 LTS
       NVIDIA-SMI： 535.113.01，
       nvcc： release release 11.8, V11.8.89，
       cudnn：8.9.7.29
       部署后，需要部署TensorRT，部署TensorRT-10.2.0.19
       onnxruntime-gpu：1.18.1，安装的适配cuda11.8的版本

运行：rembg i -m u2net dog.jpg dog-u2net.jpg
结果：命令运行可正常输出结果，运行中gpu被使用，代码中运行也正常。