环境配置

 torch version 2.0.1+cu118 Found.
 python 3.10
 onnxruntime-gpu==1.16.0
 服务器cuda 12.2 cudnn 9xxx

测试是否使用的gpu

option = onnxruntime.SessionOptions()
providers = ['CUDAExecutionProvider']
onnx_ort = onnxruntime.InferenceSession(onnx_path, sess_options=option, providers=providers)
print('================onnx_ort ', onnx_ort.get_providers())

输出结果是:==============onnx_ort [‘CUDAExecutionProvider’, ‘CPUExecutionProvider’]

onnx转换的时候不需要特殊指定,infer阶段指定CUDAExecutionProvider即可实现gpu 推理,如果没有送到gpu 上,需要修改路径

export PATH=/usr/local/cuda-11.4/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-11.4/lib64:$LD_LIBRARY_PATH
其中cuda的路径,在虚拟环境下需要修改为:/xxx/anaconda3/envs/env_name/lib
Logo

腾讯云面向开发者汇聚海量精品云计算使用和开发经验,营造开放的云计算技术生态圈。

更多推荐