Python多进程处理代码示例

模型推理时经常会用多张卡加快推理速度，通过多进程，每个进程根据gpu_index加载到不同卡上进行推理，如下为主要代码：

from multiprocessing import Process
import time
from tqdm import tqdm


def run(gpu_index, image_list):
    print(f"gpu{gpu_index}开始处理！")
    for img in tqdm(image_list, desc=f"GPU{gpu_index}", unit="张"):
        print(f"当前gpu索引{gpu_index}, 正在处理 {img}\n")
        time.sleep(1)

def get_images_list(image_path):
    from pathlib import Path
    image_files = []
    input_path = Path(image_path)
    for ext in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff'] :
        image_files.extend(input_path.glob(f"*{ext}"))
        image_files.extend(input_path.glob(f"*{ext.upper()}"))
    image_files = sorted(set(image_files))

    if not image_files:
        print(f"\n在目录 {image_path} 中未找到图片文件")
    print(f"\n找到 {len(image_files)} 张图片")

    return image_files
    
def split_images_for_gpus(image_files, num_gpus):
    gpu_images = [[] for _ in range(num_gpus)]
    for i, img_file in enumerate(image_files):
        gpu_idx = i % num_gpus
        gpu_images[gpu_idx].append(img_file)
    
    return gpu_images

def main():
    
    image_path = r"图片路径"
    image_files = get_images_list(image_path)

    gpu_nums = [0, 1, 2, 5]

    gpu_images = split_images_for_gpus(image_files, len(gpu_nums))

    processes = []
    for i in range(len(gpu_nums)):
        p = Process(target=run,
                    args=(gpu_nums[i], gpu_images[i],))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()
    print("所有进程处理完毕！")

if __name__ == "__main__":
    #multiprocessing.freeze_support()
    main()