| | |
|---|
| | |
|
|---|
| | | #include <fstream>
|
|---|
| | | #include <opencv2/core/cuda.hpp>
|
|---|
| | |
|
|---|
| | | #include "vp_infer_node.h"
|
|---|
| | |
|
|---|
| | | namespace vp_nodes {
|
|---|
| | |
|---|
| | | try {
|
|---|
| | | net = cv::dnn::readNet(model_path, model_config_path);
|
|---|
| | | #ifdef VP_WITH_CUDA
|
|---|
| | | // net.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);
|
|---|
| | | // net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA);
|
|---|
| | | // 检查可用的 GPU 数量
|
|---|
| | | int gpu_count = cv::cuda::getCudaEnabledDeviceCount();
|
|---|
| | | if (gpu_count > 0) {
|
|---|
| | | // 初始化随机数种子并随机选择一个 GPU
|
|---|
| | | std::srand(static_cast<unsigned>(std::time(nullptr)));
|
|---|
| | | int selected_gpu = std::rand() % gpu_count;
|
|---|
| | | // 设置 CUDA 设备并配置为使用指定 GPU
|
|---|
| | | cv::cuda::setDevice(selected_gpu);
|
|---|
| | | net.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);
|
|---|
| | | net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA);
|
|---|
| | |
|
|---|
| | | VP_INFO(vp_utils::string_format("[%s] Using CUDA on GPU %d", node_name.c_str(), selected_gpu));
|
|---|
| | | } else {
|
|---|
| | | VP_WARN(vp_utils::string_format("[%s] No CUDA-enabled GPUs detected. Running on CPU.", node_name.c_str()));
|
|---|
| | | }
|
|---|
| | | net.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);
|
|---|
| | | net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA);
|
|---|
| | | #endif
|
|---|
| | | }
|
|---|
| | | catch(const std::exception& e) {
|
|---|
