YOLOv8在C++调用训练时自动启动多进程致内存占满问题

我在Python中实现了基于YOLOv8的产品检测功能，通过自定义类ProductDetection可正常完成自定义数据集的训练与测试。现通过封装成wrapper.pyd在C++中调用该类，前期流程正常，但执行TrainDET启动训练时，系统自动开启大量并行进程，直至耗尽CPU内存。

Python训练函数 #

def Train(self, freeze=0, recipe_name="dummy"):
    with self.lock:
        recipe_set_path = os.path.join(self.recipe_path, str(recipe_name))
        self.model.train(data=self.data_path,
                         epochs=self.epochs,
                         time=0.5,
                         patience=10,
                         batch=self.batch_size,
                         imgsz=self.image_size,
                         cache=False,
                         save=True,  
                         resume=False,
                         name=recipe_set_path,
                         exist_ok=True,
                         device=self.device,
                         workers=0,
                         seed=42,
                         single_cls=True,
                         close_mosaic=0,
                         profile=True,
                         freeze=0,
                         plots=False)
        return

C++训练函数 #

bool TrainDET(std::string recipe_name, float freez_layer, bool cache, bool profile) {
#ifndef _DEBUG
    PyGILState_STATE gstate;
    gstate = PyGILState_Ensure();
    PyObject* pcache = cache ? Py_True : Py_False;
    PyObject* pprofile = profile ? Py_True : Py_False;

    PyObject* result = PyObject_CallMethod((PyObject*)pInstance, "Train", "(sfOO)", recipe_name.c_str(), freez_layer, pcache, pprofile);
    if (!result) {
        PyErr_Print();
        Py_DECREF(result);
        PyGILState_Release(gstate);
        return false;
    }
    Py_DECREF(result);
    PyGILState_Release(gstate);
#endif
    return true;
}

C++主函数 #

int main() {

    Py_Initialize();

    PyEval_InitThreads(); // Initialize Python's Global Interpreter Lock (GIL) for single-threaded execution

    PyInitialize(true);

    GetPyInstance(4, 2, 640, 640, "mywrapper");

    std::string recipe_model_path = "C:\\Users\\Karan\\Documents\\Visual Studio 2015\\Projects\\testwrapper\\testwrapper\\DATA\\Models\\Blisbeat_BlisterTopA\\";

    SetDETPath(recipe_model_path, "C:\\Users\\Karan\\Documents\\Visual Studio 2015\\Projects\\testwrapper\\testwrapper\\DATA\\BASE_DATASET\\data.yaml", "DATA\\DL\\det_model.pt");

    LoadDLModel("dummy");

    TrainDET("Unknown", 0, false, true);
return 0;
}

已尝试的无效解决方法 #

• 设置OMP_NUM_THREADS等环境变量为1
• 使用PyGILState_Ensure()锁定线程
• 在Python中引入threading模块

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