如何在Python类方法中并行化np.apply_along_axis？

最小可复现示例 #

我有一个MyClass类，是实际问题的抽象简化。其中run_algorithm方法调用时会生成一个形状为(100,10)的数组z，并在每次迭代中对z的每一行调用function方法。

目前使用np.apply_along_axis(self.function, 1, z)实现，但希望将这一步并行化。

存在两个复杂点：

• 在run_algorithm()的每次迭代中，self.function指向的实际方法会切换（本示例中差异较小）。
• self.function是类方法且动态变化，导致无法使用标准multiprocessing方案。

import numpy as np

class MyClass:
    
    def __init__(self, B=10, n_iterations=10):
        self.B = B 
        self.n_iterations = n_iterations
    
    def function1(self, z_row):
        """Takes numpy array z_row of dimension (d,) and returns
        numpy array of dimension (B+1, d)"""
        return np.arange(1, self.B+2).reshape(-1, 1) * z_row
    
    def function2(self, z_row):
        """Similar to function1 in terms of shapes, but does something slightly
        different."""
        return np.arange(1, (2*self.B)+2).reshape(-1, 1) * z_row
    
    def run_algorithm(self):
        """At each iteration chooses to apply either function1 or function2."""
        z = np.random.randn(100, 10)
        for i in range(self.n_iterations):
            print("Iteration: ", i)
            # Choose the function
            self.function = self.function1 if (i % 2 == 1) else self.function2
            # Apply the function using numpy (not in parallel)
            z_final = np.apply_along_axis(self.function, 1, z)
            # choose a random slice and start again
            index = np.random.choice(a=np.arange(self.B+1), size=1)[0]
            z = z_final[:, index, :]
        return z

失效方案 #

若self.function不是类方法，可通过如下multiprocessing方案实现并行：

def function_parallel(z):
    """Parallel version."""
    try:
        with Pool(8) as p:
            results = p.map(self.function, product(z))
            return results
    except KeyboardInterrupt:
        p.terminate()
    except Exception as e:
        print('Exception occurred: ', e)
        p.terminate()
    finally:
        p.join()
self.function_parallel = function_parallel

但该方案失效，因为multiprocessing要求函数定义在main中，而本场景中函数依赖类的状态。

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