【Python】使用 onnxruntime-gpu 进行推理，解决运行时间久了显存被拉爆了

原创

原小明呢 2024-03-22 22:10:09 ©著作权

文章标签 python 人工智能开发语言 ide CUDA 文章分类 Html/CSS 前端开发

©著作权归作者所有：来自51CTO博客作者原小明呢的原创作品，请联系作者获取转载授权，否则将追究法律责任

1. 问题

使用 onnxruntime-gpu 进行推理，解决运行时间久了显存被拉爆了

2. C++/Python 配置

运行时，配置 provder , gpu_mem_limit 参数来进行限制，比如2G显存

2147483648
2 * 1024 * 1024 * 1024

Python

providers = [
            (
                "TensorrtExecutionProvider",
                {
                    "device_id": 0,
                    "trt_max_workspace_size": 2147483648,
                    "trt_fp16_enable": True,
                },
            ),
            (
                "CUDAExecutionProvider",
                {
                    "device_id": 0,
                    "arena_extend_strategy": "kNextPowerOfTwo",
                    "gpu_mem_limit": 2 * 1024 * 1024 * 1024,
                    "cudnn_conv_algo_search": "EXHAUSTIVE",
                    "do_copy_in_default_stream": True,
                },
            ),
        ]

如运行时，使用 cuda 进行推理

self.session = onnxruntime.InferenceSession(
            path_or_bytes=model_file,
            providers=[
                (
                    "CUDAExecutionProvider",
                    {
                        "device_id": 0,
                        "arena_extend_strategy": "kNextPowerOfTwo",
                        "gpu_mem_limit": 2 * 1024 * 1024 * 1024,
                        "cudnn_conv_algo_search": "EXHAUSTIVE",
                        "do_copy_in_default_stream": True,
                    },
                )
            ],
        )

C++

OrtSessionOptions* session_options = /* ... */;

OrtCUDAProviderOptions options;
options.device_id = 0;
options.arena_extend_strategy = 0;
options.gpu_mem_limit = 2 * 1024 * 1024 * 1024;
options.cudnn_conv_algo_search = OrtCudnnConvAlgoSearchExhaustive;
options.do_copy_in_default_stream = 1;

SessionOptionsAppendExecutionProvider_CUDA(session_options, &options);