如何在 Python 中使用 Multiprocessing 並行執行以下代碼

Question

我有一個 function generate(file_path) ，它返回一個 integer 索引和一個 numpy 數組。 生成function的簡化如下：

def generate(file_path):
  temp = np.load(file_path)
  #get index from the string file_path
  idx = int(file_path.split["_"][0])
  #do some mathematical operation on temp
  result = operate(temp)
  return idx, result

我需要遍歷一個目錄並將generate(file_path)的結果收集到一個 hdf5 文件中。 我的序列化代碼如下：

for path in glob.glob(directory):
    idx, result = generate(path)

    hdf5_file["results"][idx,:] = result
    
hdf5_file.close()

我希望寫一個多線程或多進程的代碼來加速上面的代碼。 我該如何修改它？ 非常感謝！

我的嘗試是修改我的生成 function 並修改我的“主要”，如下所示：

def generate(file_path):
    temp = np.load(file_path)
    #get index from the string file_path
    idx = int(file_path.split["_"][0])
    #do some mathematical operation on temp
    result = operate(temp)
      
    hdf5_path = "./result.hdf5"
    hdf5_file = h5py.File(hdf5_path, 'w')
    hdf5_file["results"][idx,:] = result

    hdf5_file.close()

if __name__ == '__main__':
    ##construct hdf5 file
    hdf5_path = "./output.hdf5"
    hdf5_file = h5py.File(hdf5_path, 'w')
    hdf5_file.create_dataset("results", [2000,15000], np.uint8)

    hdf5_file.close()

    path_ = "./compute/*"
    p = Pool(mp.cpu_count())
    p.map(generate, glob.glob(path_))
    hdf5_file.close()
   
    print("finished")

但是，它不起作用。 它會拋出錯誤

KeyError: "Unable to open object (object 'results' doesn't exist)"

Answer 1

在檢查您的代碼后，我在初始化數據集時發現了一些錯誤；

您在生成 function 中生成了路徑為“./result.hdf5”的hdf5 文件。

但是，我認為您忽略了在該文件下創建一個“結果”數據集，因為這就是導致Object Does Not Exist 問題的原因。

如果您仍然面臨與錯誤消息相同的問題，請回復

Answer 2

您可以使用線程或進程池同時執行多個 function 調用。 這是一個使用進程池的示例：

from concurrent.futures import ProcessPoolExecutor
from time import sleep


def generate(file_path: str) -> int:
    sleep(1.0)
    return file_path.split("_")[1]


def main():
    file_paths = ["path_1", "path_2", "path_3"]
    
    with ProcessPoolExecutor() as pool:
        results = pool.map(generate, file_paths)
        
        for result in results:
            # Write to the HDF5 file
            print(result)
    

if __name__ == "__main__":
    main()

請注意，您不應同時寫入同一個 HDF5 文件，即文件寫入不應generate function 中的 append。