numba库加速python代码,占用内存很多

@jit(nopython=True)  
def itertools_combinations(pool, r):
    n = len(pool)
    indices = list(range(r))
    empty = not(n and (0 < r <= n))

    if not empty:
        result = [pool[i] for i in indices]
        yield result

    while not empty:
        i = r - 1
        while i >= 0 and indices[i] == i + n - r:
            i -= 1
        if i < 0:
            empty = True
        else:
            indices[i] += 1
            for j in range(i+1, r):
                indices[j] = indices[j-1] + 1

            result = [pool[i] for i in indices]
            yield result

下面这一段是主要的内容,调用corr_feature_map函数,@jit(nopython=True)能够极大加速代码,但是运行过后就会增加很多内存;如果去掉@jit(nopython=True),运行速度减慢很多,但是内存不增加。

@jit(nopython=True)
def corr_feature_map(input_single, kernel_size, strides):
    row_width = input_single.shape[0]
    column_width = input_single.shape[1]
    column_slices = int((column_width-kernel_size)/strides)

    iteration = list(itertools_combinations(np.arange(row_width), 2)) 
    feature_map = []
    for i,j in iteration:
        corr_list = []
        for k in range(column_slices+1):
            data_1 = input_single[i, k*strides:k*strides+kernel_size]
            data_2 = input_single[j, k*strides:k*strides+kernel_size]
            if np.max(data_1)==np.min(data_1) or np.max(data_2)==np.min(data_2):  # 避免(a, a, a,......, a)的方差(=0)位于分母
                corr_coef = 0
            else:
                corr_coef = np.corrcoef(data_1, data_2)[0, 1]
            corr_list.append(corr_coef)
        feature_map.append(corr_list)

    return feature_map

这是很正常的,因为它要把所有的依赖库都用jit编译才能调用,你用了np,这个库挺大的。
用空间换时间了。