尝试用Cython优化NumPy向量化运算性能未获提升的原因咨询

我尝试通过Cython化（或验证可行性）来加速一个基于NumPy的向量化运算。该代码根据两个距离矩阵（target_distances和由扁平化坐标向量计算得到的map_distances）以及距离类型信息（取值范围0-3，本次测试设为全0）计算某种应力值。测试后Cython版本与NumPy版本耗时相近（约4.47秒 vs 4.62秒），无法理解原因：NumPy和SciPy是否在后台进行了并行计算？查看核心使用率并未发现并行，通过设置export MKL_NUM_THREADS=1等环境变量关闭并行后也无变化。是我忽略了向量化运算的并行机制，还是这些库的例程已优化到极致？

NumPy版本代码（pycalculus.py） #

import scipy
import numpy as np
from scipy.spatial import distance

def decay(x, s=10):
    return scipy.special.expit(s*x)

def stress(z, target_distances, dim, distance_types, step,
           nrows, ncols):
    row_coords = np.reshape(z[:dim*nrows],(nrows,dim))
    col_coords = np.reshape(z[dim*nrows:dim*(nrows+ncols)],(ncols,dim))
    
    map_distances = distance.cdist(row_coords, col_coords).copy()
    error = target_distances - map_distances
    
    I0 = (distance_types==0) | (distance_types==3)
    I2 = distance_types == 2

    return np.sum(error[I0]**2) + np.sum((error[I2] + step)**2*decay(error[I2] + step))

Cython化版本代码（calculus.pyx） #

import cython
cimport cython
from libc.stdlib cimport malloc, free

cdef extern from "ctools.h":
    double stress (double*, double**, int**, double, int, int, int)

@cython.boundscheck(False) 
@cython.wraparound(False)
@cython.nonecheck(False)     
def stress_cython(double[:] z, double[:,:] target_distances, int [:,:]
           distance_types, double step, int dim, int nrows, int ncols):
    
    cdef int i
    cdef int j
    cdef int N = (nrows+ncols)*dim
    
    z_C = <double*>malloc(sizeof(double)*N)
    target_distances_C = <double **>malloc(sizeof(double*)*nrows)
    distance_types_C = <int **>malloc(sizeof(int*)*nrows)
    
    for i in range(N):
        z_C[i] = z[i]
    
    for i in range(nrows):
        target_distances_C[i] = <double *>malloc(sizeof(double)*ncols)
        distance_types_C[i] = <int *>malloc(sizeof(int)*ncols)
        
        for j in range(ncols):
            target_distances_C[i][j] = target_distances[i,j]
            distance_types_C[i][j] = distance_types[i,j]
        
   
    stress_val = stress(z_C, target_distances_C, 
                        distance_types_C, step, nrows, ncols, dim)
    
    for i in range(nrows):
        free(target_distances_C[i])
        free(distance_types_C[i])
        
    free(target_distances_C)
    free(distance_types_C)
    
    return stress_val

外部C实现文件（ctools.c） #

#include<stdio.h>
#include<math.h>

double decay(double x){
    return 1/(1+exp(-10*x));
}

double** dist_pairs(double* z, int nrows, int ncols, int dims)
{
    int i,j,d;
    double coord1, coord2;
    double** dist;
    
    dist = (double **)malloc(sizeof(double*)*nrows);
    for (i=0; i<nrows; i++){
        dist[i] = (double *)malloc(sizeof(double)*ncols);
    }
    
    for (i=0; i<nrows; i++){
       for (j=0; j<ncols; j++){
            dist[i][j] = 0;
            for (d=0; d<dims; d++){
                coord1 = z[d*(nrows+ncols) + i];
                coord2 = z[d*(nrows+ncols) + nrows + j];
                dist[i][j] += pow(coord1-coord2,2.0) ;
            }
            dist[i][j] = sqrt(dist[i][j]);
        }
    }
                   
    return dist;
}

double stress(double* z, double **target_distances, int** distance_types, 
              double step, int nrows, int ncols, int dim){
    int i,j;
    double stress = 0.0;
    double err = 0.0;
    double **map_distances = dist_pairs(z, nrows, ncols, dim);

    for (i=0; i<nrows; i++){
        for (j=0; j<ncols; j++){
            if (distance_types[i][j]==0 || distance_types[i][j]==3){
                stress += pow(target_distances[i][j] - map_distances[i][j],2.0);
            }
            else if (distance_types[i][j]==2){
                err = target_distances[i][j] - map_distances[i][j] + step;
                stress += pow(step,2)*decay(step);
            }
        }
    }

    for (i=0; i<nrows; i++){
        free(map_distances[i]);
    }
    free(map_distances);

    return stress;
}

测试代码（test.py） #

import numpy as np
import time 
from calculus import stress_cython
from pycalculus import stress

nrows = 3000
ncols = 2000
dim = 2
N = 100
is_discrete = 1.0

dt0 = 0
dt1 = 0
difs = []
for i in range(N):
    coordinates = np.random.rand(dim, nrows+ncols)
    coordinates_flat = coordinates.flatten()
    
    target_distances = np.random.rand(nrows, ncols)
    distance_types = np.zeros((nrows,ncols), dtype='i')
    
    t0 = time.time()
    stress1 = stress_cython(coordinates_flat, target_distances, distance_types, is_discrete, dim, nrows,
                   ncols)
    t1 = time.time()
    
    stress2 = stress(coordinates.T.flatten(), target_distances, dim, distance_types, is_discrete,
                     nrows,ncols)
    
    t2 = time.time()
    
    dt0 += t1-t0
    dt1 += t2-t1
    
    difs.append(stress1-stress2)

print(f'cython:{dt0:.2f} python:{dt1:.2f}')

补充说明 #

• 编译Cython代码时使用了-ffast-math参数
• 最终会用所有核心处理不同初始条件，因此NumPy自动并行对当前场景无帮助，这也是想弄清原因的重要因素

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