將多處理用於for循環，Python

Question

我有一個for循環，它使用一些二進制條件，並最終相應地寫入文件。 我的問題是，許多文件的條件都是正確的（有時需要寫入約1000個文件）。 因此，編寫它們會花費很長時間（大約10分鍾）。 我知道我可以以某種方式使用Python的multiprocessing並利用某些核心。

這是有效的代碼，但僅使用一個內核。

for i,n in enumerate(halo_param.strip()):
    mask = var1['halo_id'] == n
    newtbdata = tbdata1[mask]
    hdu = pyfits.BinTableHDU(newtbdata)
    hdu.writeto(('/home/Documments/file_{0}.fits').format(i))

我發現可以使用來自multiprocessing Pool來完成。

if __name__ == '__main__': pool = Pool(processes=4)

我想知道如何做到這一點，並充分利用我的至少4個核心。

Answer 1

將for循環體重組為一個函數，並將Pool.map與該函數一起使用。

def work(arg):
    i, n = arg
    mask = var1['halo_id'] == n
    newtbdata = tbdata1[mask]
    hdu = pyfits.BinTableHDU(newtbdata)
    hdu.writeto(('/home/Documments/file_{0}.fits').format(i))

if __name__ == '__main__':
    pool = Pool(processes=4)
    pool.map(work, enumerate(halo_param.strip()))
    pool.close()
    pool.join()