ONNX 转换后的 TensorFlow 保存模型在 CPU 上运行，但不在 GPU 上运行

如何解决ONNX 转换后的 TensorFlow 保存模型在 CPU 上运行，但不在 GPU 上运行

我使用 tf2onnx 将 tensorflow 保存的模型转换为 ONNX 格式：

python3 -m tf2onnx.convert --saved-model saved_model/ --output onnx/model.onnx --opset 11

转换工作正常，我可以使用 cpu 对 ONNX 模型进行推理。

我安装了 onnxruntime-gpu 以使用 GPU 运行推理并遇到错误：

RuntimeException: [ONNXRuntimeError] : 6 : RUNTIME_EXCEPTION : Non-zero status code returned while running Relu node. Name:'FirstStageFeatureExtractor/resnet_v1_101/resnet_v1_101/conv1/Relu' Status Message: /onnxruntime_src/onnxruntime/core/providers/cuda/cuda_call.cc:97 bool onnxruntime::CudaCall(ERRTYPE,const char*,ERRTYPE,const char*) [with ERRTYPE = cudaError; bool THRW = true] /onnxruntime_src/onnxruntime/core/providers/cuda/cuda_call.cc:91 bool onnxruntime::CudaCall(ERRTYPE,const char*) [with ERRTYPE = cudaError; bool THRW = true] CUDA failure 2: out of memory ; GPU=0 ; hostname=coincoin; expr=cudamalloc((void**)&p,size); 
Stacktrace:

Stacktrace:

我是唯一使用 Titan RTX（24GB RAM）GPU 的人。该模型使用其 tensorflow 保存的模型版本在 GPU 上运行良好，具有 10GB 的 GPU 内存。

版本是：

• 张量流 1.14.0
• CUDA 10.0
• CuDNN 7.6.5
• onnx 1.6.0
• onnxruntime 1.1.0
• tf2onnx 1.9.2
• python 3.6
• Ubuntu 18.04

解决方法

根据您的信息和版本，可能有两种解决方案：

• 1. 信息内存不足，请检查gpu内存是否可用，如果不可用，设置gpu内存：

config = tf.ConfigProto()                                                                           
config.gpu_options.visible_device_list = "0"                                                        
config.gpu_options.per_process_gpu_memory_fraction = 0.1                                            
set_session(tf.Session(config=config)) 

• 2. 降级 onnxruntime-gpu，根据你的版本，CUDA 10.0 匹配 onnxruntime-gpu==1.0.0，而不是 1.1.0，需要 cuda 11 https://github.com/Microsoft/onnxruntime/releases/tag/v1.1.0