Python多处理池;等待迭代完成

时间:2018-09-25 19:08:23

标签: python numpy multiprocessing

我有一个很大的数据集,我希望我的脚本进行迭代,在每个条目上执行一系列操作,然后排列结果以存储到HDD。由于数据集可能相对较大(〜250 GB),因此RAM的可用性要求数据集必须一次处理1000个条目的块(在下面的代码中我称之为dataBlock)。我还使用multiprocessing.Pool类来促进使用多个CPU内核来完成此任务。

我基本上已经安排好了一切,以便将每个dataBlock传递到Pool,Pool使用imap方法对dataBlock执行所需的计算,Pool返回计算结果,数据块将附加到列表。此列表(processed_data)是计算集所需的最终产品。

processed_data = []

multiprocessing.Pool(processor_cap) as pool:

    for blockIndex, block in enumerate(range(1000, height-remainder, 1000)):

        #Read-in 1000 spectra from source dataset
        dataBlock = np.asarray(raw_dset[blockIndex*1000:block][:])

        '''
        Pass data block to processor pool, which iterates through data
        block. Each spectrum is handed off to a CPU in the pool,
        which centroids it and appends the result to "processed_block".
        '''
        processed_block = pool.imap(centroid_spectrum, dataBlock)

        #Append processed spectra to processed data bin
        for idx, processed_spectrum in enumerate(processed_block):
            processed_data.append(processed_spectrum)

我想知道的是如何在调用pool.imap()之后使脚本暂停,直到返回完整的processed_block而不关闭池为止。当前,它直接进入for循环,该循环立即在上面的代码段中,而无需等待processed_block返回pool.imap。我试过在pool.join()调用之后立即调用pool.imap(),但是它只返回***AssertionError,并再次继续到它下面的for循环。一旦将所有dataBlocks都馈入池中,就在上面最外面的pool.close()循环的末尾,我最终可以在脚本中稍后成功地调用pool.join()for

提前感谢您的帮助!

2 个答案:

答案 0 :(得分:0)

如果不付出很多努力来改变周围的情况,那么很难处理您的示例;但是如果您从imap()调用中获得一个Iterator,则可以考虑在到达for循环之前将Iterator的元素解析为一个列表:

processed_block = pool.imap(centroid_spectrum, dataBlock)
processed_block = [ x for x in processed_block ] # convert from an iterator to a list
for idx, processed_spectrum in enumerate(processed_block):

达到您想要的目标吗?

答案 1 :(得分:0)

我只是将Pool.imap()调用更改为Pool.map()调用,并且脚本按预期运行。有关更多信息,请参见与Mikhail Burshteyn的交流。