机器学习入门：消息队列 数据预处理

最近给一个入门项目补数据管道，一开始我图省事，写了个脚本扫目录直接预处理，结果数据一多就卡在那儿，终端半天不动，我还以为代码写挂了。后来才发现，不是模型的问题，也不是 Python 的问题，是我把“读文件、解码图片、缩放、归一化、落盘”全塞一个循环里了，前面还没处理完，后面又不停往里灌。这个场景其实很适合把消息队列拉进来：生产端只管把“要处理哪张图”丢进队列，消费端慢慢做预处理，谁也别催谁。

先把 RabbitMQ 跑起来

我这里用 RabbitMQ，入门够用，管理界面也直观。首先把这个 docker 镜像拉下来：

docker pull docker.1ms.run/rabbitmq:3-management

我这里配了加速域名，不需要的或者拉的时候报错了，可以把 docker.1ms.run/ 删除，没准儿你用的时候这个加速域名已经不可用了。然后启动：

docker run -d --name mq \
  -p 5672:5672 \
  -p 15672:15672 \
  docker.1ms.run/rabbitmq:3-management

不出问题的话，浏览器打开 http://localhost:15672，用户名密码都是 guest。能看到管理界面就说明队列服务起来了。Python 这边装依赖：

pip install pika pillow numpy

生产者别急着发大文件

我一开始想过直接把图片二进制塞进消息里，后来想了想，这玩意儿不合适。队列不是网盘，消息太大既占内存又拖速度。比较稳的做法是消息里只放路径或者对象存储的 key，消费者拿到路径自己去读。

下面这个生产者会把 raw_images 目录下所有 jpg 扔进 raw_images 队列：

import json
from pathlib import Path
import pika

connection = pika.BlockingConnection(
    pika.ConnectionParameters(host="localhost")
)
channel = connection.channel()

channel.queue_declare(queue="raw_images", durable=True)

for img_path in Path("raw_images").glob("*.jpg"):
    message = json.dumps({"image_path": str(img_path.resolve())})
    channel.basic_publish(
        exchange="",
        routing_key="raw_images",
        body=message,
        properties=pika.BasicProperties(delivery_mode=2),
    )
    print("sent", img_path)

connection.close()

delivery_mode=2 是让消息持久化，队列也声明成 durable=True。这样 RabbitMQ 重启后，至少不会全丢。当然，真正生产环境还得考虑磁盘和确认机制，入门阶段先别把坑挖太大。

消费者里做预处理

消费者这边就干一件事：收消息，打开图片，转 RGB，缩到 224x224，归一化，存成 .npy，然后手动 ack。这里我把 prefetch_count 设成 1，避免一次拉太多消息把内存顶上去。

import json
from pathlib import Path

import numpy as np
import pika
from PIL import Image

connection = pika.BlockingConnection(
    pika.ConnectionParameters(host="localhost")
)
channel = connection.channel()
channel.queue_declare(queue="raw_images", durable=True)
channel.queue_declare(queue="preprocessed", durable=True)
channel.basic_qos(prefetch_count=1)

output_dir = Path("preprocessed")
output_dir.mkdir(exist_ok=True)

def callback(ch, method, properties, body):
    data = json.loads(body)
    src = Path(data["image_path"])
    try:
        img = Image.open(src).convert("RGB")
        img = img.resize((224, 224))
        arr = np.asarray(img, dtype=np.float32) / 255.0

        out_path = output_dir / f"{src.stem}.npy"
        np.save(out_path, arr)

        ch.basic_publish(
            exchange="",
            routing_key="preprocessed",
            body=json.dumps({"npy_path": str(out_path.resolve())}),
            properties=pika.BasicProperties(delivery_mode=2),
        )
        ch.basic_ack(delivery_tag=method.delivery_tag)
        print("done", out_path)
    except Exception as e:
        print("failed", src, e)
        # 这里别无限 requeue，坏文件会一直转圈
        ch.basic_nack(delivery_tag=method.delivery_tag, requeue=False)

channel.basic_consume(
    queue="raw_images",
    on_message_callback=callback,
    auto_ack=False,
)

print("waiting...")
channel.start_consuming()

跑的时候先开消费者，再跑生产者。这样你能看到消息被一条条消费，预处理结果落到 preprocessed 目录，同时 preprocessed 队列里也有后续训练脚本可以接的消息。训练脚本想批量取数据，就从 preprocessed 队列里读，或者直接读目录，别和预处理搅在一起。

几个容易翻车的点

第一个是路径。生产者发的是宿主机绝对路径，消费者如果跑在另一个容器里，路径很可能对不上。要么把数据目录挂载到相同路径，要么消息里发相对路径加根目录配置，别硬编码。

第二个是 ack。用 auto_ack=True 看着省事，但消费者处理到一半挂了，消息就没了。老老实实手动 ack，出异常就 nack，但别 requeue=True 无限重试，坏图片会把你队列堵死。

第三个是别把消息队列当万能药。数据量很小、预处理几秒跑完的时候，直接写脚本更简单。消息队列真正有用的地方是解耦、削峰和失败重试，比如一边有多个采集端往里丢任务，另一边预处理速度跟不上，或者你想把预处理和训练拆成两个进程慢慢跑。入门阶段先把这条链路跑通，比一上来堆 Kafka 实在得多。