假设你有一个模型 predict_func,它可以对单条数据进行预测。我们要对大量数据 data_list 进行批量预测。
from joblib import Parallel, delayed
from tqdm import tqdm
import numpy as np
import time # 用于模拟耗时预测函数
# -----------------------------
# 模拟预测函数(可以替换为模型预测)
def predict_func(x):
time.sleep(0.01) # 模拟计算耗时
return x ** 2
# -----------------------------
# -----------------------------
# 批量预测工具
def batch_predict(data, n_jobs=4, batch_size=100):
"""
data: list 或 numpy 数组
n_jobs: 并行任务数量
batch_size: tqdm 显示进度条的批量大小
"""
results = []
# 将数据拆分成批次
batches = [data[i:i+batch_size] for i in range(0, len(data), batch_size)]
# 并行处理每个批次
def process_batch(batch):
return [predict_func(x) for x in batch]
# tqdm 包装 batches,用于显示进度条
results = Parallel(n_jobs=n_jobs)(
delayed(process_batch)(batch) for batch in tqdm(batches, desc="预测进度")
)
# 将结果扁平化
results = [item for sublist in results for item in sublist]
return results
# -----------------------------
if __name__ == "__main__":
# 模拟大数据
data_list = list(range(1000))
preds = batch_predict(data_list, n_jobs=4, batch_size=50)
print("前 10 个预测结果:", preds[:10])
🔹 说明
tqdm 显示进度条
desc="预测进度" 可以自定义显示文字
对大数据可清楚看到批量处理进度
joblib.Parallel + delayed 并行处理
n_jobs 控制并行进程数量
适合 CPU 密集型任务
批量处理可以减少任务调度开销
批次处理
通过 batch_size 将数据拆分,兼顾进度显示和并行效率
扁平化结果
Parallel 返回的是每个批次的列表,需要用 [item for sublist in results for item in sublist]
来自 “ ITPUB博客 ” ,链接:https://blog.itpub.net/70050070/viewspace-3117474/ ,如需转载,请注明出处,否则将追究法律责任。