Python C扩展中正确内存释放问题求助

问题背景 #

你实现了生成延迟向量的C函数delay_vectors，并编写Python包装器将其暴露给NumPy。循环调用该函数时内存持续增长，即使尝试设置NPY_ARRAY_OWNDATA标记也无法解决内存泄漏问题。

原C函数实现 #

double** delay_vectors(const double* ts, int L, int tau, int p) {
    int i, j;

    // Allocate contiguous memory for v_i_list
    double **v_i_list = (double **)malloc((L-tau*(p-1)) * sizeof(double *));
    v_i_list[0] = (double *)malloc((L-tau*(p-1)) * p * sizeof(double));
    
    for (int i = 1; i < (L-tau*(p-1)); i++) {
        v_i_list[i] = v_i_list[i - 1] + p;
    }
    

    for (i = 0; i < L - tau * (p - 1); ++i) {
        for (j = 0; j < p; ++j) {
            v_i_list[i][j] = ts[i + j * tau];
        }
    }

    return v_i_list;
}

原Python包装器实现 #

static PyObject* delay_vectors_wrapper(PyObject* self, PyObject* args) {
    PyObject *input_array;
    PyArrayObject *input_array_np;
    int tau, p;

    // Parse the input arguments
    if (!PyArg_ParseTuple(args, "Oii", &input_array, &tau, &p)) {
        return NULL;
    }

    // Convert the input object to a NumPy array
    input_array_np = (PyArrayObject *)PyArray_FROM_OTF(input_array, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);

    if (input_array_np == NULL) {
        return NULL;
    }


    int ndim = PyArray_NDIM(input_array_np);
    if (ndim != 1) {
        PyErr_SetString(PyExc_ValueError, "Input array must be 1-dimensional.");
        Py_DECREF(input_array_np);
        return NULL;
    }
    
    int L = PyArray_DIM(input_array_np, 0);

    // Get data pointer from the NumPy array
    double *ts = (double *)PyArray_DATA(input_array_np);

    double** result = delay_vectors(ts, L, tau, p);

    // Convert the C result to a np array
    npy_intp dims[2] = {L - tau * (p - 1), p};
    PyObject* output_array = PyArray_SimpleNewFromData(2, dims, NPY_DOUBLE, result);
    

    
    Py_DECREF(input_array_np);

    return output_array;
}

泄漏原因分析 #

你的delay_vectors函数分配了两块独立内存：

1. 指针数组v_i_list（存储每行的指针）
2. 连续的数据块v_i_list[0]（存储实际延迟向量数据）

而PyArray_SimpleNewFromData仅对传入的data指针（你这里错误传入了result，实际应为result[0]）进行包装。即使设置NPY_ARRAY_OWNDATA，NumPy也只会释放result[0]指向的数据块，指针数组v_i_list的内存永远不会被回收，导致每次调用都泄漏这块内存。

修复方案 #

方案1：简化内存分配（推荐） #

既然数据已经是连续存储的，完全不需要二级指针数组。修改delay_vectors直接返回单一的连续数据块，让NumPy直接管理这块内存：

修改后的C函数

double* delay_vectors(const double* ts, int L, int tau, int p) {
    int i, j;
    const int num_vectors = L - tau * (p - 1);
    // 直接分配连续的二维数据内存（按行优先存储）
    double* data = (double*)malloc(num_vectors * p * sizeof(double));
    
    if (!data) {
        // 内存分配失败时设置Python异常
        PyErr_NoMemory();
        return NULL;
    }

    for (i = 0; i < num_vectors; ++i) {
        for (j = 0; j < p; ++j) {
            data[i * p + j] = ts[i + j * tau];
        }
    }

    return data;
}

修改后的包装器

static PyObject* delay_vectors_wrapper(PyObject* self, PyObject* args) {
    PyObject *input_array;
    PyArrayObject *input_array_np;
    int tau, p;

    if (!PyArg_ParseTuple(args, "Oii", &input_array, &tau, &p)) {
        return NULL;
    }

    input_array_np = (PyArrayObject *)PyArray_FROM_OTF(input_array, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
    if (!input_array_np) {
        return NULL;
    }

    if (PyArray_NDIM(input_array_np) != 1) {
        PyErr_SetString(PyExc_ValueError, "Input array must be 1-dimensional.");
        Py_DECREF(input_array_np);
        return NULL;
    }
    
    const int L = PyArray_DIM(input_array_np, 0);
    double* ts = (double*)PyArray_DATA(input_array_np);

    const int num_vectors = L - tau * (p - 1);
    double* result_data = delay_vectors(ts, L, tau, p);
    if (!result_data) {
        Py_DECREF(input_array_np);
        return NULL;
    }

    const npy_intp dims[2] = {num_vectors, p};
    PyObject* output_array = PyArray_SimpleNewFromData(2, dims, NPY_DOUBLE, result_data);
    
    // 让NumPy接管数据内存的所有权，GC时自动释放
    PyArray_ENABLEFLAGS((PyArrayObject*)output_array, NPY_ARRAY_OWNDATA);

    Py_DECREF(input_array_np);
    return output_array;
}

方案2：添加自定义清理函数（保留原C函数） #

如果不想修改原C函数，可以通过PyCapsule创建持有指针数组的对象，并绑定到NumPy数组上，让Python GC在回收数组时同时释放两块内存：

步骤1：定义清理函数

static void delay_vectors_cleanup(void* ptr) {
    double** v_i_list = (double**)ptr;
    if (v_i_list) {
        free(v_i_list[0]); // 释放数据块
        free(v_i_list);    // 释放指针数组
    }
}

步骤2：修改包装器

static PyObject* delay_vectors_wrapper(PyObject* self, PyObject* args) {
    PyObject *input_array;
    PyArrayObject *input_array_np;
    int tau, p;

    if (!PyArg_ParseTuple(args, "Oii", &input_array, &tau, &p)) {
        return NULL;
    }

    input_array_np = (PyArrayObject *)PyArray_FROM_OTF(input_array, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
    if (!input_array_np) {
        return NULL;
    }

    if (PyArray_NDIM(input_array_np) != 1) {
        PyErr_SetString(PyExc_ValueError, "Input array must be 1-dimensional.");
        Py_DECREF(input_array_np);
        return NULL;
    }
    
    const int L = PyArray_DIM(input_array_np, 0);
    double* ts = (double*)PyArray_DATA(input_array_np);

    double** result = delay_vectors(ts, L, tau, p);
    if (!result) {
        Py_DECREF(input_array_np);
        return NULL;
    }

    const npy_intp dims[2] = {L - tau * (p - 1), p};
    // 注意这里传入的是实际数据块指针result[0]
    PyObject* output_array = PyArray_SimpleNewFromData(2, dims, NPY_DOUBLE, result[0]);
    if (!output_array) {
        free(result[0]);
        free(result);
        Py_DECREF(input_array_np);
        return NULL;
    }

    // 创建Capsule对象持有指针数组，绑定清理函数
    PyObject* capsule = PyCapsule_New(result, NULL, delay_vectors_cleanup);
    if (!capsule) {
        Py_DECREF(output_array);
        free(result[0]);
        free(result);
        Py_DECREF(input_array_np);
        return NULL;
    }

    // 将Capsule设置为NumPy数组的base对象，GC时自动触发清理
    if (PyArray_SetBaseObject((PyArrayObject*)output_array, capsule) < 0) {
        Py_DECREF(capsule);
        Py_DECREF(output_array);
        free(result[0]);
        free(result);
        Py_DECREF(input_array_np);
        return NULL;
    }

    Py_DECREF(input_array_np);
    return output_array;
}

验证 #

修改后再次运行循环测试，内存占用将不再持续增长。用valgrind检测也会显示无内存泄漏。

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