tensorrt指定使用显卡

在使用 TensorRT 进行模型推理时，确保使用 GPU 并指定具体的 GPU 设备可以通过以下步骤实现：

1. 确保 TensorRT 使用 GPU

TensorRT 会默认使用 能用的GPU 进行推理，只要你的环境中有可用的 GPU 并且你正确安装了 CUDA 和 cuDNN。用下面代码可以验证是否有gpu可用。

import tensorrt as trt

# 检查是否有可用的 GPU
def check_available_gpus():
    import pycuda.driver as cuda
    cuda.init()
    gpu_count = cuda.Device.count()
    print(f"Available GPUs: {gpu_count}")
    for i in range(gpu_count):
        gpu = cuda.Device(i)
        print(f"GPU {i}: {gpu.name()}")

check_available_gpus()

2. 设置 CUDA 环境变量

你可以通过设置 CUDA_VISIBLE_DEVICES 环境变量来指定 TensorRT 使用哪个 GPU。例如：

export CUDA_VISIBLE_DEVICES=0  # 使用第0块GPU

或者在你的 Python 代码中：

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'  # 使用第0块GPU

3.使用cudart来指定显卡

from cuda import cudart
#确保在 TensorRT 运行时设置 CUDA 设备
cudart.cudaSetDevice(0)

注意

如果不通过上面方式指定显卡，tensorrt会默认使用能用的显卡，如果指定了则使用指定的显卡