laugh12321 / TensorRT-YOLO

🚀 TensorRT-YOLO: Supports YOLOv3, YOLOv5, YOLOv6, YOLOv7, YOLOv8, YOLOv9, YOLOv10, and PP-YOLOE using TensorRT acceleration with EfficientNMS, CUDA Kernels and CUDA Graphs!
https://github.com/laugh12321/TensorRT-YOLO
GNU General Public License v3.0
383 stars 50 forks source link

[Help]: 请问如何实现指定GPU进行推理? #27

Closed leidingzhang closed 2 months ago

leidingzhang commented 2 months ago

我使用python代码测试yolov8模型,虽然可以推理成功,但是查看代码未发现指定GPU的地方

laugh12321 commented 2 months ago

这里没有写(采用的是默认的Gpu 0),如果你有多个设备需要指定可以使用cuda.cudart.cudaSetDevice自己加上。