如何在Databricks中规模化YOLOv8批量推理（处理1亿张S3图像）

针对你的场景，必须采用MLflow管理模型+Spark分布式计算+Autoloader加载数据的组合实现规模化，替代当前多独立笔记本的低效模式，具体方案如下：

一、要不要用MLflow？ #

肯定要上MLflow，它解决了规模化推理的三个核心痛点：

• 统一模型版本管理：避免不同节点加载不同模型版本的混乱
• 标准化依赖环境：确保所有分布式节点的Ultralytics、YOLOv8依赖完全一致
• 简化分布式部署：可直接将注册的MLflow模型转换成Spark UDF，快速适配集群并行计算

二、分步实现流程 #

1. 将YOLOv8模型注册到MLflow #

自定义PyFunc适配YOLOv8的推理逻辑，打包模型到MLflow仓库：

import mlflow
import mlflow.pyfunc
from ultralytics import YOLO
from pyspark.sql.functions import col, explode

class YOLOv8Model(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        # 加载本地训练好的YOLOv8模型
        self.model = YOLO(context.artifacts["model_path"])
    
    def predict(self, context, model_input):
        # 接收图像路径列表，返回结构化检测结果
        results = self.model(model_input.iloc[:, 0].tolist())
        output = []
        for path, res in zip(model_input.iloc[:, 0], results):
            for box in res.boxes:
                output.append({
                    "image_path": path,
                    "x1": box.xyxy[0][0].item(),
                    "y1": box.xyxy[0][1].item(),
                    "x2": box.xyxy[0][2].item(),
                    "y2": box.xyxy[0][3].item(),
                    "class": res.names[int(box.cls)],
                    "confidence": box.conf.item()
                })
        return output

# 打包并注册模型
with mlflow.start_run():
    mlflow.pyfunc.log_model(
        artifact_path="yolov8-model",
        python_model=YOLOv8Model(),
        artifacts={"model_path": "/dbfs/path/to/your/yolov8/model"},
        conda_env={
            "channels": ["conda-forge"],
            "dependencies": [
                "python=3.9",
                "pip",
                {"pip": ["ultralytics==8.0.200", "mlflow"]}
            ]
        }
    )
model_version = mlflow.register_model("runs:/<your-run-id>/yolov8-model", "yolov8-production")

2. 用Autoloader批量加载S3图像 #

Autoloader针对大规模云存储优化，支持自动扫描、增量读取，避免重复处理：

# 全量读取S3中的图像，生成包含路径的DataFrame
image_df = (spark.readStream
            .format("cloudFiles")
            .option("cloudFiles.format", "image")
            .option("cloudFiles.includeExistingFiles", "true")
            .option("cloudFiles.maxFilesPerTrigger", 10000)  # 单次处理文件数，按需调整
            .load("s3://your-bucket/image-directory/")
            .select("path"))

3. 分布式GPU推理：MLflow模型转Spark UDF #

将MLflow模型转换成Spark UDF，让每个GPU executor并行处理数据：

# 加载注册的模型为Spark UDF
yolov8_udf = mlflow.pyfunc.spark_udf(spark, f"models:/yolov8-production/{model_version.version}")

# 应用UDF并展开检测结果
results_df = image_df.select(
    "path",
    yolov8_udf("path").alias("detections")
).withColumn("detection", explode("detections"))

# 提取结构化字段
final_results_df = results_df.select(
    "path",
    col("detection.x1"),
    col("detection.y1"),
    col("detection.x2"),
    col("detection.y2"),
    col("detection.class"),
    col("detection.confidence")
)

4. 结果持久化到Delta Lake #

将推理结果存入Delta Lake，支持高效查询、ACID事务和后续分析：

(final_results_df.writeStream
 .format("delta")
 .option("checkpointLocation", "dbfs:/checkpoints/yolov8-inference/")
 .option("path", "dbfs:/delta/yolov8-inference-results/")
 .trigger(availableNow=True)  # 一次性处理全量数据，完成后自动停止
 .start()
 .awaitTermination())

三、关键优化点 #

• 集群配置：使用GPU实例（如AWS g4dn.2xlarge），每个executor分配1个GPU，设置参数：

  spark.executor.resource.gpu.amount=1
  spark.task.resource.gpu.amount=1
  spark.executor.instances=20  # 按需扩容，提升并行度
  

• 分区调整：根据图像大小调整Spark分区数，确保每个分区包含300-500张图像，最大化GPU利用率
• 增量推理：后续新增图像时，移除includeExistingFiles参数，自动处理新增文件
• 资源调度：使用Databricks按需集群或Serverless集群，任务完成后自动释放资源，降低成本

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