Python 的 C 扩展中的内存泄漏

Memory Leak in C-extension for Python

这是我第一次为 python 编写 C 扩展,您可以看到我丑陋且可能超级低效的卷积 C++ 实现。我的内存管理有问题。每次我在 python 中调用此函数时,它都会消耗大约 500MB 的内存(对于大小为 100x112x112x3 的批处理和大小为 3x3x3x64 的内核)并且之后不会释放它。即使这不是 class 方法,我是否真的要注意引用计数?或者我是否必须在代码的某处手动释放内存?请注意,我排除了所有错误检查以获得更好的概览。谢谢

PyObject* conv2d(PyObject*, PyObject* args)

{
    PyObject* data;
    PyObject* shape;
    PyObject* kernel;
    PyObject* k_shape;
    int stride;

    PyArg_ParseTuple(args, "OOOOi", &data, &shape, &kernel, &k_shape, &stride);

    Py_ssize_t dims = PyTuple_Size(shape);
    Py_ssize_t kernel_dims = PyTuple_Size(k_shape);

    int shape_c[3];
    int k_shape_c[4];

    for (int i = 0; i < kernel_dims; i++)
    {
        if (i < dims)
        {
            shape_c[i] = PyLong_AsLong(PyTuple_GetItem(shape, i));
        }
        k_shape_c[i] = PyLong_AsLong(PyTuple_GetItem(k_shape, i));
    }

    PyObject* data_item, kernel_item;
    PyObject* ret_array = PyList_New(0);
    double conv_val, channel_sum;

    for (int oc = 0; oc < k_shape_c[3]; oc++)
    {
        for (int row = 0; row < shape_c[0]; row += stride)
        {
            for (int col = 0; col < shape_c[1]; col += stride)
            {
                channel_sum = 0;
                for (int ic = 0; ic < shape_c[2]; ic++)
                {
                    conv_val = 0;
                    for (int k_row = 0; k_row < k_shape_c[0]; k_row++)
                    {
                        for (int k_col = 0; k_col < k_shape_c[1]; k_col++)
                        {
                            data_item = PyList_GetItem(data, row + k_row);
                            if (!data_item)
                            {
                                PyErr_Format(PyExc_IndexError, "Index out of bounds");
                                return NULL;
                            }
                            data_item = PyList_GetItem(data_item, col + k_col);
                            data_item = PyList_GetItem(data_item, ic);
                            kernel_item = PyList_GetItem(kernel, k_row);
                            kernel_item = PyList_GetItem(kernel_item, k_col);
                            kernel_item = PyList_GetItem(kernel_item, ic);
                            kernel_item = PyList_GetItem(kernel_item, oc);
                            conv_val += PyFloat_AsDouble(data_item) * PyFloat_AsDouble(kernel_item);
                        }
                    }
                    channel_sum += conv_val;
                }
                PyList_Append(ret_array, PyFloat_FromDouble(channel_sum));
            }
        }
    }
    return ret_array;
}

泄漏来自:

PyList_Append(ret_array, PyFloat_FromDouble(channel_sum));

PyFloat_FromDouble 创建一个新的引用,PyList_Append 获得引用的共享所有权(它不会 窃取 / 消耗 参考)。当使用 PyList_Append 并且您希望 list 拥有您自己的引用时,您必须在附加后明确释放您的引用,例如(省略错误检查):

PyObject *pychannel_sum = PyFloat_FromDouble(channel_sum);
PyList_Append(ret_array, pychannel_sum);
Py_DECREF(pychannel_sum);

另一种(如果合适的话速度更快)解决方案是将 list 预分配到正确的大小,并用 PyList_SetItem/PyList_SET_ITEM 填充条目,两者 窃取一个引用,而不是增加引用计数。一般来说,没有明确提及引用窃取的 API 不会,您需要监管自己的引用计数。

请注意 memory-wise,单个 PyFloat 比 C double(它们包装)贵很多;在 64 位系统上,list 中的每个 PyFloat 占用 32 个字节(list 中的指针占用 8 个字节,PyFloat 本身占用 24 个字节),而 PyFloat 本身占用 8 个字节原始 C double.

您可能需要研究使用 Python's array module(创建正确 size/type 的 array,使用缓冲协议对其进行 C 级视图,然后填写缓冲区);代码会稍微复杂一些,但内存使用量会下降 4 倍。 numpy 类型将提供相同的优势(并且结果可以更灵活地使用)。