如何在 numpy 布尔索引中防止这种内存泄漏

Question

我正在尝试优化 numpy 中的一种方法，评估图像中许多位置的高斯函数。现在我正在计算每个高斯中心位置周围的布尔掩码，并使用它们来索引数据数组，以节省计算时间。对于少量数据点，与通过np.argwhere 计算索引相比，布尔索引大大加快了脚本的速度。

不幸的是，布尔索引似乎会导致内存泄漏。 以下示例导致我的 RAM 达到最大值。我该如何防止这种情况？这是一个麻木的错误，还是我做错了什么？我认为这可能是因为布尔索引创建了数据的副本，正如有人在this answer 中所写的那样。不再需要副本后不应该释放内存吗？有没有办法手动释放内存？

import numpy as np
import matplotlib.pyplot as plt


class Image():
    def __init__(self, nx, ny, px):
        x = np.linspace(0, nx * px, nx)
        y = np.linspace(0, ny * px, ny)
        self.x, self.y = np.meshgrid(x, y)
        self.data = np.zeros((nx, ny), dtype=np.uint8)
        self.nx = nx  # number of pixels
        self.ny = ny
        self.px = px  # pixel size [mm]

    def gaussian(self, x, y, x0, y0, amplitude=None, sigma=0.01):
        """ 2d gaussian function

        Parameters:
            x, y: coordinates of evaluation position
            x0, y0: center coordinates of Gaussian
            amplitude: amplitude of the Gaussian
            sigma: width of Gaussian
        """
        if amplitude is None:
            amplitude = np.iinfo(self.data.dtype).max
        result = amplitude * np.exp(
            -((x - x0)**2 + (y - y0)**2) / (2 * sigma ** 2))
        return result.astype(self.data.dtype)

    def mask_from_pos(self, x0, y0, radius_px):
        """ returns a square mask around the position (x0, y0)
        """
        masks = np.zeros((x0.size, self.ny, self.nx), dtype=bool)
        i_center_x = (x0 // self.px).astype(int)
        i_center_y = (y0 // self.px).astype(int)
        for ix, iy, mask in zip(i_center_x, i_center_y, masks):
            y_start = max(0, min(iy - radius_px, self.ny - 1))
            y_end = max(0, min(iy + radius_px + 1, self.ny))
            x_start = max(0, min(ix - radius_px, self.nx - 1))
            x_end = max(0, min((ix + radius_px + 1), self.nx))

            mask[y_start: y_end, x_start: x_end] = True
        return masks

    def add_gaussians(self, posx, posy, radius_px=9):
        masks = self.mask_from_pos(posx, posy, radius_px)
        for x0, y0, mask in zip(posx, posy, masks):
            self.data[mask] += self.gaussian(
                self.x[mask], self.y[mask], x0, y0).astype(self.data.dtype)


if __name__ == '__main__':
    image = Image(2000, 2000, 0.005)
    xy_max = 2000 * 0.005
    x0 = np.random.rand(3000) * xy_max
    y0 = np.random.rand(3000) * xy_max
    image.add_gaussians(x0, y0)

    plt.figure(dpi=300, figsize=(8, 8))
    plt.imshow(image.data,
            cmap=plt.get_cmap('gray'),
            extent=[0, xy_max, 0, xy_max])

我已经尝试过调试，但似乎没有任何 python 对象增长，所以我认为这是 numpy 中的一个问题，但我不是 100% 确定。非常感谢您对此的任何帮助！提前致谢。

编辑 1：解决方法

我找到了一个解决方案，它给出了相同的结果，但不会溢出内存。只需用掩码生成器替换掩码数组的计算就可以了：

    def mask_from_pos_gen(self, x0, y0, radius_px):
        """ yields a square mask around the position (x0, y0)
        """
        i_center_x = (x0 // self.px).astype(int)
        i_center_y = (y0 // self.px).astype(int)
        for ix, iy in zip(i_center_x, i_center_y):
            mask = np.zeros(self.data.shape, dtype=bool)
            y_start = max(0, min(iy - radius_px, self.ny - 1))
            y_end = max(0, min(iy + radius_px + 1, self.ny))
            x_start = max(0, min(ix - radius_px, self.nx - 1))
            x_end = max(0, min((ix + radius_px + 1), self.nx))

            mask[y_start: y_end, x_start: x_end] = True
            yield mask

然后可以使用与之前的 mask_from_pos 完全相同的方式使用此方法，但通过这种方式释放内存。把这个留在这里，以防其他人有同样的问题，并希望一些 numpy Guru 能解释这种行为。

Answer 1

事实证明，布尔索引并不是解决此问题的最佳解决方案。与布尔掩码相比，使用 numpy 切片可大幅提升速度。

    def slice_from_pos_gen(self, x0, y0, radius_px):
        """ yields a square mask around the position (x0, y0)
        """
        i_center_x = (x0 // self.px).astype(int)
        i_center_y = (y0 // self.px).astype(int)
        for ix, iy in zip(i_center_x, i_center_y):
            y_start = max(0, min(iy - radius_px, self.ny - 1))
            y_end = max(0, min(iy + radius_px + 1, self.ny))
            x_start = max(0, min(ix - radius_px, self.nx - 1))
            x_end = max(0, min((ix + radius_px + 1), self.nx))

            yield np.s_[y_start: y_end, x_start: x_end]

这个方法可以和之前的 mask_from_pos 完全一样的方式使用，但是速度更快并且不会填满内存。