Colab中如何用subprocess实现YOLOv4的GPU并发批量推理

原有逐张循环调用Darknet检测命令的流程，存在大量无效开销：每次启动独立进程都要重新加载YOLOv4权重、初始化CUDA上下文，单张推理时Tesla V100的显存、算力利用率通常不足10%，大规模测试集下90%的时间都消耗在重复的进程初始化上，而非实际推理。

优先方案：Darknet原生批量推理（零额外并发开销） #

绝大多数支持CUDA编译的Darknet版本，都原生支持通过路径列表批量处理图片，无需自行管理并发，仅需启动一次进程、加载一次权重即可跑完所有测试图，速度是逐张调用的10~20倍。

实现步骤：

• 收集所有测试目录下的PNG图片绝对路径，写入临时列表文件
• 单次调用Darknet命令传入列表参数，批量完成所有推理
• 代码示例：

import os
import glob
import subprocess

# 测试图目录配置
TEST_IMG_DIR = "/content/darknet/test_images"
# 收集所有PNG路径
test_imgs = sorted(glob.glob(os.path.join(TEST_IMG_DIR, "*.png")))

# 生成Darknet可读取的路径列表文件
with open("test_imgs_list.txt", "w") as f:
    f.write("\n".join(test_imgs))

# 单次调用Darknet完成全量推理
darknet_cmd = [
    "./darknet", "detector", "test",
    "cfg/obj.data", "cfg/yolov4-custom.cfg", "backup/yolov4-custom_last.weights",
    "-list", "test_imgs_list.txt",
    "-thresh", "0.5",
    "-ext_output", # 输出检测框坐标
    "-dont_show", # 跳过可视化渲染，无需要保存检测图时加这个参数提速30%
    "-out", "detection_results.json" # 结果统一输出到json
]
subprocess.run(darknet_cmd, check=True)

备选方案：Subprocess受控并发（适配不支持列表参数的旧版Darknet） #

如果使用的Darknet版本没有原生-list参数，可通过拆分任务、多进程并发的方式提升利用率，Tesla V100 16G显存下单YOLOv4 416输入的推理进程仅占用_{1.5G显存，可同时开3}4个并发进程不会触发OOM。

实现步骤：

• 拆分全量测试图为多个小批次，每个批次对应一个独立的Darknet子进程
• 用进程池管理并发数量，避免显存过载
• 代码示例：

import os
import glob
import subprocess
from concurrent.futures import ProcessPoolExecutor

# 环境适配参数
TEST_IMG_DIR = "/content/darknet/test_images"
BATCH_SIZE_PER_WORKER = 20 # 单个子进程一次处理20张图，减少循环调用
MAX_WORKERS = 3 # 16G V100适配3个并发，32G显存可调整为5

# 拆分全量列表为小批次
def split_list(lst, batch_size):
    for i in range(0, len(lst), batch_size):
        yield lst[i:i+batch_size]

# 单个子进程的推理任务
def run_darknet_single_batch(img_batch):
    tmp_list = f"/tmp/batch_{os.getpid()}.txt"
    with open(tmp_list, "w") as f:
        f.write("\n".join(img_batch))
    
    cmd = [
        "./darknet", "detector", "test",
        "cfg/obj.data", "cfg/yolov4-custom.cfg", "backup/yolov4-custom_last.weights",
        tmp_list, # 旧版Darknet直接传入单批次的路径列表
        "-thresh", "0.5",
        "-dont_show",
        "-out", f"/tmp/result_{os.getpid()}.json"
    ]
    subprocess.run(cmd, check=True, capture_output=True)
    os.remove(tmp_list)
    return f"/tmp/result_{os.getpid()}.json"

if __name__ == "__main__":
    # 收集所有测试图
    test_imgs = sorted(glob.glob(os.path.join(TEST_IMG_DIR, "*.png")))
    # 拆分为子批次
    img_batches = list(split_list(test_imgs, BATCH_SIZE_PER_WORKER))
    # 并发执行所有批次
    with ProcessPoolExecutor(max_workers=MAX_WORKERS) as executor:
        result_files = list(executor.map(run_darknet_single_batch, img_batches))
    # 最后合并所有结果文件即可

适配Colab V100的注意事项 #

• 并发数不要盲目调高：跑1张图时用!nvidia-smi查看单进程Darknet的显存占用，按单进程占用的1.5倍预留CUDA上下文开销计算最大并发数，避免Colab OOM杀进程
• 不要用多线程启动子进程：Python GIL会限制线程并发执行，多进程才能真正调用CUDA多流资源
• 所有文件存在Colab本地磁盘：测试图、权重、Darknet代码都放在/content目录下，不要读取挂载的Google Drive路径，避免IO瓶颈拖慢速度

内容的提问来源于stack exchange，提问作者mchd