python - 如何在 cython 的 nogil 循环下将 numpy.ndarray 分配给临时变量？

Question

我正在尝试实现 implicit推荐模型，并且在代码运行时计算 ~11kk 用户超过 ~100k 项目的前 5 个建议时存在问题。

我能够通过 numpy 和一些 cython sparkles(在 jupyter notebook 中)部分解决问题。使用 numpy 排序的行仍然使用单核:

%%cython -f
# cython: language_level=3
# cython: boundscheck=False
# cython: wraparound=False
# cython: linetrace=True
# cython: binding=True
# distutils: define_macros=CYTHON_TRACE_NOGIL=1
from cython.parallel import parallel, prange
import numpy as np
from tqdm import tqdm

def test(users_items=np.random.rand(11402139//1000, 134751//100)
        , int N=5, show_progress=True, int num_threads=1):
    # Define User count and loops indexes
    cdef int users_c = users_items.shape[0], u, i
    # Predefine zero 2-D C-ordered array for recommendations
    cdef int[:,::1] users_recs = np.zeros((users_c, N), dtype=np.intc)
    for u in tqdm(range(users_c), total=users_c, disable=not show_progress):
        # numpy .dot multiplication using multiple cores
        scores = np.random.rand(134751//1000, 10).dot(np.random.rand(10))
        # numpy partial sort
        ids_partial = np.argpartition(scores, -N)[-N:]
        ids_top = ids_partial[np.argsort(scores[ids_partial])]
        # Fill predefined 2-D array
        for i in range(N):
            users_recs[u, i] = ids_top[i]
    return np.asarray(users_recs)
# Working example
tmp = test()

我对其进行了分析 - np.argpartition 消耗了 60% 的函数时间并使用了 onde 核心。我试图让它并行，因为我有一个 80 核的服务器。因此，我对一部分用户(使用多核)执行 .dot 操作，并计划通过 numpy 排序结果(使用单核)并行填充空的预定义数组，但我遇到了问题标题中的错误:

%%cython -f
# cython: language_level=3
# cython: boundscheck=False
# cython: wraparound=False
# cython: linetrace=True
# cython: binding=True
# distutils: define_macros=CYTHON_TRACE_NOGIL=1
from cython.parallel import parallel, prange
import numpy as np
from tqdm import tqdm
from math import ceil
def test(int N=10, show_progress=True, int num_threads=1):
    # Define User and Item count and loops indexes
    cdef int users_c = 11402139//1000, items_c = 134751//100, u, i, u_b
    # Predefine zero 2-D C-ordered array for recommendations
    cdef int[:,::1] users_recs = np.zeros((users_c, N), dtype=np.intc)
    # Define memoryview var
    cdef float[:,::1] users_items_scores_mv
    progress = tqdm(total=users_c, disable=not show_progress)
    # For a batch of Users
    for u_b in range(5):
        # Use .dot operation which use multiple cores
        users_items_scores = np.random.rand(num_threads, 10).dot(np.random.rand(134751//100, 10).T)
        # Create memory view to 2-D array, which I'm trying to sort row wise
        users_items_scores_mv = users_items_scores
        # Here it starts, try to use numpy sorting in parallel
        for u in prange(num_threads, nogil=True, num_threads=num_threads):
            ids_partial = np.argpartition(users_items_scores_mv[u], items_c-N)[items_c-N:]
            ids_top = ids_partial[np.argsort(users_items_scores_mv[u][ids_partial])]
            # Fill predefined 2-D array
            for i in range(N):
                users_recs[u_b + u, i] = ids_top[i]
        progress.update(num_threads)
    progress.close()
    return np.asarray(users_recs)

得到这个(full error):

Error compiling Cython file:
------------------------------------------------------------
...
        # Create memory view to 2-D array,
        # which I'm trying to sort row wise
        users_items_scores_mv = users_items_scores
        # Here it starts, try to use numpy sorting in parallel
        for u in prange(num_threads, nogil=True, num_threads=num_threads):
            ids_partial = np.argpartition(users_items_scores_mv[u], items_c-N)[items_c-N:]
           ^
------------------------------------------------------------

/datascc/enn/.cache/ipython/cython/_cython_magic_201b296cd5a34240b4c0c6ed3e58de7c.pyx:31:12: Assignment of Python object not allowed without gil

我阅读了有关内存 View 和分配分配的信息，但没有找到适用于我的情况的示例。

Answer 1

我最终得到了自定义 C++ 函数，它通过 openmp 与 nogil 并行填充 numpy 数组。它需要用 cython 重写 numpy 的 argpartition 部分排序。算法是这样的(3-4可以循环):

1. 定义空数组 A[i,j] 和 memory view B_mv[i,k];其中“i”是批量大小，“j”是一些列，“k”是排序后要返回的所需项目数
2. 在 A&B 的内存上创建指针
3. 运行一些计算并用数据填充 A
4. 并行遍历 i-s 并填充 B
5. 将结果转换为可读形式

解决方案包括:

topnc.h - 自定义函数实现的头文件:

/* "Copyright [2019] <Tych0n>"  [legal/copyright] */
#ifndef IMPLICIT_TOPNC_H_
#define IMPLICIT_TOPNC_H_

extern void fargsort_c(float A[], int n_row, int m_row, int m_cols, int ktop, int B[]);

#endif  // IMPLICIT_TOPNC_H_

topnc.cpp - 函数体:

#include <vector>
#include <limits>
#include <algorithm>
#include <iostream>

#include "topnc.h"

struct target {int index; float value;};
bool targets_compare(target t_i, target t_j) { return (t_i.value > t_j.value); }

void fargsort_c ( float A[], int n_row, int m_row, int m_cols, int ktop, int B[] ) {
    std::vector<target> targets;
    for ( int j = 0; j < m_cols; j++ ) {
        target c;
        c.index = j;
        c.value = A[(n_row*m_cols) + j];
        targets.push_back(c);
    }
    std::partial_sort( targets.begin(), targets.begin() + ktop, targets.end(), targets_compare );
    std::sort( targets.begin(), targets.begin() + ktop, targets_compare );
    for ( int j = 0; j < ktop; j++ ) {
        B[(m_row*ktop) + j] = targets[j].index;
    }
}

ctools.pyx - 用法示例

# distutils: language = c++
# cython: language_level=3
# cython: boundscheck=False
# cython: wraparound=False
# cython: nonecheck=False
from cython.parallel import parallel, prange
import numpy as np
cimport numpy as np

cdef extern from "topnc.h":
    cdef void fargsort_c ( float A[], int n_row, int m_row, int m_cols, int ktop, int B[] ) nogil

A = np.zeros((1000, 100), dtype=np.float32)
A[:] = np.random.rand(1000, 100).astype(np.float32)
cdef:
    float[:,::1] A_mv = A
    float* A_mv_p = &A_mv[0,0]
    int[:,::1] B_mv = np.zeros((1000, 5), dtype=np.intc)
    int* B_mv_p = &B_mv[0,0]
    int i
for i in prange(1000, nogil=True, num_threads=10, schedule='dynamic'):
    fargsort_c(A_mv_p, i, i, 100, 5, B_mv_p)
B = np.asarray(B_mv)

compile.py——编译文件；在终端中通过命令“python compile.py build_ext --inplace -f”运行它(这将生成文件 ctools.cpython-*.so，然后用于导入):

from os import path
import numpy
from setuptools import setup, Extension
from Cython.Distutils import build_ext
from Cython.Build import cythonize

ext_utils = Extension(
    'ctools'
    , sources=['ctools.pyx', 'topnc.cpp']
    , include_dirs=[numpy.get_include()]
    , extra_compile_args=['-std=c++0x', '-Os', '-fopenmp']
    , extra_link_args=['-fopenmp']
    , language='c++'
)

setup(
    name='ctools',
    setup_requires=[
        'setuptools>=18.0'
        , 'cython'
        , 'numpy'
    ]
    , cmdclass={'build_ext': build_ext}
    , ext_modules=cythonize([ext_utils]),
)

它用于 adding “推荐所有”功能到 implicit肌萎缩侧索硬化模型。