3、Batch推理的工程实现：数据加载与预处理、模型前向传播的Batch处理、输出结果的解Batch与后处理

好，咱们今天聊聊Batch推理在工程上到底怎么落地。说实话，很多同学在论文里看Batch推理觉得挺简单——不就是把一堆数据塞进去一起算吗？但真到了写代码的时候，坑一个接一个。我当年第一次做线上推理服务，就因为在数据加载这步没处理好，导致GPU利用率跟心电图似的，忽高忽低。

这一章，我会把整个流程拆成三段来讲：数据怎么喂进去、模型怎么算、结果怎么吐出来。你跟着我走一遍，基本就能避开我踩过的那些坑了。

3.1 数据加载与预处理：别让CPU拖了GPU的后腿

先问个问题：你的GPU在等数据的时候在干嘛？答案是——闲着。这是最常见的性能杀手。

我个人习惯用PyTorch的DataLoader，但很多人只是简单调个接口，没理解背后的机制。咱们来看一个典型的配置：

from torch.utils.data import DataLoader, Dataset

class MyDataset(Dataset):
    def __init__(self, data_list):
        self.data = data_list

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        # 这里做单个样本的加载和预处理
        sample = self.data[idx]
        # 假设是图片，需要解码、resize、归一化
        image = decode_and_preprocess(sample)
        return image

# 关键参数在这里
dataloader = DataLoader(
    MyDataset(all_samples),
    batch_size=32,
    shuffle=False,
    num_workers=4,      # 多进程加载
    pin_memory=True,    # 锁页内存，加速CPU到GPU传输
    prefetch_factor=2   # 预取批次数量
)

这里有几个点我要特别强调：

• num_workers：我建议设成CPU核心数的一半左右。设太多反而会因为进程切换开销导致性能下降。我在一个项目里试过设成16，结果CPU被打满，数据加载反而变慢了。
• pin_memory=True：这个一定要开。它让CPU内存页不被换出，GPU可以直接通过DMA访问，传输速度能快好几倍。
• prefetch_factor：告诉DataLoader提前准备几个batch。GPU在算第N个batch的时候，CPU已经在准备第N+1、N+2个了。

预处理这块，我建议能离线做的就别在线做。比如图片的resize、归一化参数，完全可以提前算好存起来。在线只做必要的、跟样本相关的操作。

3.2 模型前向传播的Batch处理：形状一致是铁律

数据准备好了，接下来就是模型推理。Batch处理的核心就一句话：同一批次的数据，形状必须完全一致。

为什么会这样？因为GPU的并行计算是基于张量的。你可以把GPU想象成一个超级整齐的军队，每个士兵（计算单元）做同样的事。如果数据形状不一样，军队就乱了。

实际场景中，最头疼的是变长数据。比如NLP里的句子，有的长有的短。怎么办？我常用的方法是padding：

import torch
from torch.nn.utils.rnn import pad_sequence

# 假设我们有3个样本，长度分别是3、5、4
sequences = [
    torch.tensor([1, 2, 3]),
    torch.tensor([4, 5, 6, 7, 8]),
    torch.tensor([9, 10, 11, 12])
]

# 填充到相同长度（用0填充）
padded = pad_sequence(sequences, batch_first=True, padding_value=0)
# 结果形状: [3, 5]
# 内容:
# [[1, 2, 3, 0, 0],
#  [4, 5, 6, 7, 8],
#  [9, 10, 11, 12, 0]]

# 记录实际长度
lengths = torch.tensor([3, 5, 4])

# 前向传播
with torch.no_grad():
    output = model(padded)

这里有个坑：padding之后，模型可能会对填充的位置也进行计算，产生无意义的输出。所以很多模型（比如Transformer）需要传入attention mask，告诉模型哪些位置是真实数据，哪些是填充的。

另外，batch size不是越大越好。我见过有人为了追求吞吐量，把batch size设到256甚至512。结果呢？显存爆了，或者单个batch的延迟太高，反而影响了整体性能。我的建议是：从32开始，逐步往上调，直到显存占用达到80%左右。留点余量给临时变量和碎片。

3.3 输出结果的解Batch与后处理：把张量变回有意义的东西

模型算完了，输出是一个大张量。但用户要的不是张量，是具体的预测结果。这一步叫解Batch，说白了就是把批处理的结果拆开，还原成单个样本的输出。

举个例子，假设我们做图像分类，模型输出形状是[32, 1000]（32个样本，1000个类别）。解Batch就是把这个张量拆成32个独立的预测向量：

# 假设output形状是 [batch_size, num_classes]
output = model(padded)  # 形状: [32, 1000]

# 解Batch：沿着batch维度拆开
predictions = []
for i in range(output.size(0)):
    single_output = output[i]  # 形状: [1000]
    # 后处理：取概率最大的类别
    predicted_class = torch.argmax(single_output).item()
    predictions.append(predicted_class)

# 或者用列表推导式，更Pythonic
predictions = [torch.argmax(output[i]).item() for i in range(output.size(0))]

后处理这一步，不同任务差别很大：

• 分类任务：通常就是argmax + softmax（如果需要概率值）
• 目标检测：需要NMS（非极大值抑制）去掉重复框
• NLP生成：需要解码策略，比如beam search或者采样
• 语义分割：需要把像素级的类别映射回原图尺寸

我个人习惯把后处理也写成batch版本，而不是逐个样本循环。比如NMS，很多框架（如torchvision）已经支持batch输入了：

from torchvision.ops import nms

# 假设boxes形状是 [N, 4]，scores形状是 [N]
# 对所有样本一起做NMS？不行，不同样本的框不能混在一起
# 正确做法：每个样本独立做NMS

# 更高效的方式：用batched_nms（如果框架支持）
# 或者自己写一个循环，但用向量化操作加速

3.4 整体流程可视化

说了这么多，咱们用一张图把整个流程串起来。这张图是我自己画的一个标准Batch推理流水线：

这张图里，三个阶段是串行的，但实际运行时是流水线并行的。你想想看，GPU在算第N个batch的时候，CPU已经在加载第N+1个batch的数据了。这就是为什么我说数据加载是性能瓶颈的第一嫌疑人。

3.5 总结一下

Batch推理的工程实现，说白了就是三件事：

1. 数据加载要快——用多进程、预取、锁页内存，别让GPU闲着
2. 模型计算要稳——保证batch内形状一致，该加mask加mask，batch size要调优
3. 结果处理要准——解Batch别搞错维度，后处理要跟模型输出对齐

嗯，这里要注意：不要为了Batch而Batch。如果业务场景里请求量不大，或者单个样本的计算量已经很小了，强行Batch反而会增加延迟。我见过有人把batch size设到64，结果每个请求都要等64个凑齐了才开始算，单个请求的延迟从10ms变成了500ms。这就得不偿失了。

好了，这一章的内容就到这里。代码示例都是可以直接跑的，你可以在自己的项目里试试看。记住，纸上得来终觉浅，绝知此事要躬行。

想要更多资料，联系可可