PyTorch推理扩展实战：用Ray Data轻松实现多机多卡并行

单机 PyTorch 模型跑推理没什么问题但数据量一旦上到万级、百万级瓶颈就暴露出来了内存不够、GPU 利用率低、I/O 拖后腿更别说还要考虑容错和多机扩展。传统做法是自己写多线程 DataLoader、管理批次队列、手动调度 GPU 资源这哥工程量可不小调试起来也麻烦。Ray Data 提供了一个更轻量的方案在几乎不改动原有 PyTorch 代码的前提下把单机推理扩展成分布式 pipeline。原始的 PyTorch 代码典型的推理场景模型加载、预处理、批量预测一套下来大概长这样import torch import torchvision from PIL import Image from typing import List class TorchPredictor: def __init__(self, model: torchvision.models, weights: torchvision.models): self.weights weights self.model model(weightsweights) self.model.eval() self.transform weights.transforms() self.device cuda if torch.cuda.is_available() else cpu self.model.to(self.device) def predict_batch(self, batch: List[Image.Image]) - torch.Tensor: with torch.inference_mode(): batch torch.stack([ self.transform(img.convert(RGB)) for img in batch ]).to(self.device) logits self.model(batch) probs torch.nn.functional.softmax(logits, dim1) return probs处理几张图片完全没问题predictor TorchPredictor( torchvision.models.resnet152, torchvision.models.ResNet152_Weights.DEFAULT ) images [ Image.open(/content/corn.png).convert(RGB), Image.open(/content/corn.png).convert(RGB) ] predictions predictor.predict_batch(images)大数据量图片数量从几张变成几万张、几百万张情况完全不一样了。内存撑不住不可能把所有图一股脑塞进去GPU 利用率上不去多卡场景下吞吐量优化是个棘手的问题万一跑到一半挂了怎么办分布式部署能不能用上集群资源还有个容易被忽视的点数据加载的 I/O 往往才是真正的瓶颈。自己从头写一套健壮的 pipeline 处理这些问题少说得折腾好几天。Ray Data 的思路Ray Data 是个分布式数据处理框架跟 PyTorch 配合得很好。关键是改造成本极低原有代码基本不用大动。第一步改造 Predictor 类把predict_batch方法换成__call__输入从 PIL Image 列表改成包含 numpy 数组的字典import numpy as np from typing import Dict class TorchPredictor: def __init__(self, model: torchvision.models, weights: torchvision.models): self.weights weights self.model model(weightsweights) self.model.eval() self.transform weights.transforms() self.device cuda if torch.cuda.is_available() else cpu self.model.to(self.device) def __call__(self, batch: Dict[str, np.ndarray]): Ray Data passes a dict batch with numpy arrays. # Convert numpy arrays back to PIL Images images [Image.fromarray(img_array) for img_array in batch[image]] with torch.inference_mode(): tensor_batch torch.stack([ self.transform(img.convert(RGB)) for img in images ]).to(self.device) logits self.model(tensor_batch) probs torch.nn.functional.softmax(logits, dim1) # Get top prediction top_probs, top_indices torch.max(probs, dim1) return { predicted_class_idx: top_indices.cpu().numpy(), confidence: top_probs.cpu().numpy() }改动点说明__call__替代predict_batch输入类型从List[Image.Image]变成Dict[str, np.ndarray]方法内部把 numpy 数组转回 PIL Image输出改成 dict 格式结果要搬回 CPU数据在进程间的移动由 Ray 负责。还有个细节要注意Ray Data 用 numpy 数组而非 PIL Image因为 numpy 数组跨进程序列化效率更高。第二步构建 Ray Dataset根据场景选择合适的创建方式小数据集直接从内存构建import ray import numpy as np ray.init() # Convert PIL Images to numpy arrays images [ Image.open(/path/to/image1.png).convert(RGB), Image.open(/path/to/image2.png).convert(RGB) ] # Create Ray Dataset from numpy arrays ds ray.data.from_items([{image: np.array(img)} for img in images])中等规模数据集推荐从文件路径延迟加载# Create dataset from paths image_paths [/path/to/img1.png, /path/to/img2.png] ds_paths ray.data.from_items([{path: path} for path in image_paths]) # Load images lazily def load_image(batch): images [np.array(Image.open(path).convert(RGB)) for path in batch[path]] return {image: images} ds ds_paths.map_batches(load_image, batch_size10)生产环境首选read_images()Ray 全权接管# Most efficient - Ray handles everything ds ray.data.read_images(/path/to/image/directory/) # or with specific files ds ray.data.read_images([/path/img1.png, /path/img2.png])第三步跑分布式推理核心代码如下weights torchvision.models.ResNet152_Weights.DEFAULT # Distributed batch inference results_ds ds.map_batches( TorchPredictor, fn_constructor_args(torchvision.models.resnet152, weights), batch_size32, num_gpus1, computeray.data.ActorPoolStrategy(size4) # 4 parallel actors ) # Collect results results results_ds.take_all() # Process results for result in results: class_idx result[predicted_class_idx] confidence result[confidence] print(fPredicted: {weights.meta[categories][class_idx]} ({confidence:.2%}))搞定了。新版 Ray 里concurrency参数已经废弃要换成computeActorPoolStrategy(sizeN)这种写法。改动总结自动分批Ray 自己决定最优 batch size分布式执行多 worker 并行跑GPU 调度自动把卡分配给 worker流式处理数据在 pipeline 里流动不用一次性全加载进内存容错机制worker 挂了会自动重试。生产环境RAY还可以直接读云存储的数据S3、GCS、Azure Blob 都支持# Read directly from S3, GCS, or Azure Blob ds ray.data.read_images(s3://my-bucket/images/) results ds.map_batches( predictor, batch_size64, num_gpus1, concurrency8 # 8 parallel GPU workers )多节点集群也可以用同一套代码10 台机器还是 100 台机器根本不用改# Connect to your Ray cluster ray.init(ray://my-cluster-head:10001) # Same code as before ds ray.data.read_images(s3://my-bucket/million-images/) results ds.map_batches(predictor, batch_size64, num_gpus1)进阶用法每个 batch 都重新加载模型太浪费了用 ActorPoolStrategy 让模型实例常驻内存from ray.data import ActorPoolStrategy results ds.map_batches( TorchPredictor, fn_constructor_args(torchvision.models.resnet152, weights), batch_size32, num_gpus1, computeActorPoolStrategy(size4) # Keep 4 actors alive )这样吞吐量提升很明显。CPU、GPU 资源可以细调results ds.map_batches( TorchPredictor, fn_constructor_args(torchvision.models.resnet152, weights), batch_size32, num_gpus1, # 1 GPU per actor num_cpus4, # 4 CPUs per GPU worker computeActorPoolStrategy(size8) )推理完直接写到云存储results.write_parquet(s3://my-bucket/predictions/)几个容易踩的坑Ray Data 没法直接序列化 PIL Image 对象得先转成 numpy 数组# ❌ This will fail ds ray.data.from_items([{image: pil_image}]) # ✅ This works ds ray.data.from_items([{image: np.array(pil_image)}]) # ✅ Or use read_images() (best) ds ray.data.read_images(/path/to/images/)Ray 2.51 之后concurrency不能用了# ❌ Deprecated ds.map_batches(predictor, concurrency4) # ✅ New way ds.map_batches(predictor, computeActorPoolStrategy(size4))batch size 太大容易 OOM保守起见可以从小的开始试# Monitor GPU memory and adjust batch_size accordingly results ds.map_batches( predictor, batch_size16, # Start conservative num_gpus1 )实践建议batch size 可以从小往大试观察 GPU 显存占用# Too small: underutilized GPU batch_size4 # Too large: OOM errors batch_size256 # Just right: depends on your model and GPU # For ResNet152 on a single GPU, 32-64 works well batch_size32ActorPoolStrategy 处理 20 张图大概要 9.7 秒而原生 PyTorch 跑 2 张图几乎瞬间完成。所以图片量少的时候 Ray Data 的启动开销反而不划算所以这个方案是几百上千张图的场景才能体现优势。Ray 自带 dashboard默认在 8265 端口# Check Ray dashboard at http://localhost:8265 ray.init(dashboard_host0.0.0.0)代码中可以包一层 try-except 防止单个样本出错拖垮整个任务def safe_predictor(batch: dict): try: return predictor(batch) except Exception as e: return {error: str(e), probs: None}跑之前加个计时可以进行性能 profilingimport time start time.time() results ds.map_batches(predictor, batch_size32) results.take_all() print(fProcessed in {time.time() - start:.2f} seconds)总结适合的场景数据集太大内存放不下需要多卡或多机并行长时间任务需要容错不想自己写分布式代码。不太必要的场景图片量在百张以内数据集轻松塞进内存只有一张卡而且短期内不打算扩展。Ray Data 的好处在于迁移成本低。PyTorch 代码改动很小换个方法签名、把数据包成 Ray Dataset就能换来从单卡到多机的无痛扩展、自动 batching 和并行优化、内置容错、云存储无缝对接等功能。如果你下次写多线程 data loader 或者手动管理 GPU pool 之前可以先考虑一下这哥方法把分布式系统的脏活累活交给 Ray精力留给构建模型本身。