tf2onnx转换的ONNX模型在Python中仅占用CPU运行问题

模型转换操作 #

将DeepSort仓库中的TensorFlow模型转换为ONNX格式，执行命令：

python3 -m tf2onnx.convert --graphdef  mars-small128.pb --output mars-small128_nchw.onnx --inputs-as-nchw "images:0" --inputs "images:0" --outputs-as-nchw "features:0" --outputs "features:0"

ONNX模型加载代码 #

使用以下代码加载转换后的ONNX模型：

class ImageEncoder(object):

    def __init__(self, checkpoint_filename):
        self.image_shape=[128, 64]
        
        self.session = ort.InferenceSession(checkpoint_filename, providers=['CUDAExecutionProvider'])
        print(f"the result of get_provider is: {self.session.get_providers()}")
    def __call__(self, data_x):
        out = self.session.run(["features:0"], {'images:0':data_x})[0]
        return out

异常现象 #

• 控制台输出显示CUDA执行提供者已被加载：the result of get_provider is: ['CUDAExecutionProvider', 'CPUExecutionProvider']
• 运行过程中CPU所有核心负载拉满至100%，但通过nvidia-smi查看GPU负载极低，未发挥GPU加速作用

环境版本信息 #

>>> onnx.__version__
'1.11.0'
>>> onnxruntime.__version__
'1.10.0'
>>> tf2onnx.__version__
'1.12.1'

• CUDA版本：cuda_11.0

内容的提问来源于stack exchange，提问作者Dmitry