在 python 中使用多处理优化大型数组的处理

Posted 2023-02-25

【中文标题】在 python 中使用多处理优化大型数组的处理

【英文标题】：Optimize Processing of large array using multiprocessing in python

【发布时间】：2015-12-28 09:32:19

【问题描述】：

我有一个大型数据数组（大约百万行），其中每一行都包含一个数量的坐标。我在 scipy.spatial 中使用 KDtree 算法将这些坐标分组到预定义的单元格中并附加单元格编号。在每一行的末尾。我正在使用 mulitprocessing.pool 来处理一片数组。没有。切片数等于没有。系统的核心。处理代码大约需要一分钟。我可以通过一些修改来提高速度还是我已经达到了极限。我必须处理大约 900 个这样的数组，所以时间在这里非常重要；）

我的代码相关部分如下，

  #function called by external code
  def AnodeAssignment(sim_data_dir,filename,energy_data_dir):
      argv = sim_data_dir + filename
      global data
      data = np.loadtxt(argv,str,usecols=(0,2,5,6,8,9)) #load the array from file
      good  = np.where(data[:,2]=='phot')
      data = data[good]
      data = np.delete(data,2,1).astype(float)
      slice_width = data.shape[0]/multiprocessing.cpu_count()
      #slice array based on no. of core
      data_slice = []
      for i in range(multiprocessing.cpu_count()):
         data_slice.append(data[i*slice_width:(i+1)*slice_width])
      data_slice = np.array(data_slice)
      #pool processes 
      pool = multiprocessing.Pool()
      results = pool.map(pool_process,data_slice)
      pool.close()
      data = results[0]
      for i in range(1,multiprocessing.cpu_count()): #combine results
         data = np.vstack((data,results[i]))

  #function to process each slice using KDtree algo.
  def pool_process(data_slice):
      y = np.linspace(-19.5, 19.5, 14 ,endpoint=True)
      z = np.linspace(-9, 6, 6, endpoint=True)
      yv, zv = np.meshgrid(y,z)
      tree = spatial.KDTree(zip(yv.ravel(),zv.ravel()))
      dist, indices = tree.query(data_slice[:,3:])
      zz = tree.data[indices][:,1]
      yy = tree.data[indices][:,0]

      anode = np.ndarray((data_slice[:,0].size,1),int)

      a1  = np.where(np.logical_and(np.in1d(yy,y[1:-1:2]),np.in1d(zz ,[6])))
      a2  = np.where(np.logical_and(np.in1d(yy,y[2:-1:2]),np.in1d(zz ,[6])))
      a3  = np.where(np.logical_and(np.in1d(yy,y[1:-1:2]),np.in1d(zz ,[3])))
      a4  = np.where(np.logical_and(np.in1d(yy,y[2:-1:2]),np.in1d(zz ,[3])))
      a5  = np.where(zz==0)
      a6  = np.where(zz==-3)
      a7  = np.where(zz==-6)
      a8  = np.where(yy==-19.5)
      a9  = np.where(zz==-9)
      a10 = np.where(yy==19.5)

      anode[a1]  = 1
      anode[a2]  = 2
      anode[a3]  = 3
      anode[a4]  = 4
      anode[a5]  = 5
      anode[a6]  = 6
      anode[a7]  = 7
      anode[a8]  = 8
      anode[a9]  = 9
      anode[a10] = 10

      data_slice = np.hstack((data_slice,anode))
      return data_slice

【问题讨论】：

瓶颈是什么？另外，它是 Python 2 还是 3？如果这是 Python 2，您是否考虑过使用 itertools.izip 而不是 zip？

我不确定瓶颈...这就是我想知道的是否有人可以指出...使用 python 2.7...itertools.izip 比 zip 更好吗？跨度>

尝试使用hotshot 找出瓶颈。 This question 大约是 zip 与 izip。

但无论如何，我敢打赌瓶颈在spatial.KDTree 行（即使我不是一个numpy专家）。与其学习如何使用 hotshot，不如尝试插入两个 print 语句：一个在调用 KDTree 的上方和一个下方（例如：`print "start"; tree = ...; print "stop"）并检查如果 KDTree 需要一分钟才能完成

好的，谢谢...我会尝试打印方法...

【参考方案1】：

如果可以并行进行计算，使用Pool 是一个很好的步骤。 “蛮力”类别的一个明显改进是借用具有更多内核的机器。 :-) 以同样的方式，您也许可以使用机器集群。

在优化之前，您必须测量。在您的代码上使用分析器来查看它在哪里花费时间。然后看看你是否可以做一些不同的事情来让它更快。

话虽如此，通常最大的改进通常是通过使用另一种算法来发现的。但由于我不知道您在这里要完成什么，因此无法提供具体建议。