如何实现循环真实推理调用以提升Hugging Face模型GPU利用率？

核心原理 #

数据访问瓶颈的本质是单张图推理时，GPU等待CPU加载/传输数据的时间远长于计算时间，导致GPU利用率极低。通过批量循环推理，让GPU持续处于计算状态，用计算时间掩盖数据加载的延迟，从而提升整体利用率和推理速度。

具体实现步骤及代码 #

1. 环境与依赖准备

确保安装必要库：

pip install torch transformers pillow

2. 高效批量数据加载

用PyTorch DataLoader 实现多进程数据加载，避免单张加载的延迟：

import torch
from transformers import DetrImageProcessor, DetrForObjectDetection
from PIL import Image
import os
from torch.utils.data import DataLoader, Dataset

class ImageDataset(Dataset):
    def __init__(self, image_dir, processor):
        self.image_paths = [os.path.join(image_dir, f) for f in os.listdir(image_dir) if f.endswith(('png', 'jpg'))]
        self.processor = processor

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

    def __getitem__(self, idx):
        # 加载并预处理单张图
        image = Image.open(self.image_paths[idx]).convert("RGB")
        encoding = self.processor(images=image, return_tensors="pt")
        # 移除batch维度，由DataLoader统一添加
        return {k: v.squeeze() for k, v in encoding.items()}

3. 核心循环推理逻辑

# 初始化模型与处理器，移至GPU
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50").to("cuda")
model.eval()  # 切换推理模式

# 配置DataLoader：batch_size根据GPU显存调整，num_workers设为CPU核心数的1-2倍
dataset = ImageDataset("你的图片目录路径", processor)
dataloader = DataLoader(
    dataset,
    batch_size=8,  # 示例值，可根据显存调整为16/32
    shuffle=False,
    num_workers=4,
    persistent_workers=True  # 保持worker进程，提升后续批次加载速度
)

# GPU预热：避免第一次推理因CUDA初始化导致的延迟
with torch.no_grad():
    dummy_img = Image.new("RGB", (640, 640))
    dummy_input = processor(images=dummy_img, return_tensors="pt").to("cuda")
    model(**dummy_input)

# 批量循环推理
inference_results = []
with torch.no_grad():  # 关闭梯度计算，减少内存占用
    for batch in dataloader:
        # 批量数据移至GPU
        batch = {k: v.to("cuda") for k, v in batch.items()}
        # 连续执行推理，GPU持续计算，掩盖数据加载延迟
        outputs = model(**batch)
        
        # 异步处理推理结果（不阻塞推理循环）
        target_sizes = torch.tensor([[img.shape[0], img.shape[1]] for img in batch["pixel_values"]])
        processed_outputs = processor.post_process_object_detection(
            outputs,
            target_sizes=target_sizes,
            threshold=0.5  # 置信度阈值
        )
        inference_results.extend(processed_outputs)

关键优化要点 #

• 调整batch_size：在GPU显存允许的范围内尽量调大（比如16、32），最大化GPU计算资源占用。
• num_workers优化：设为CPU核心数的1-2倍，充分利用CPU资源加载数据，避免拖慢GPU。
• 关闭梯度计算：必须使用torch.no_grad()，否则会占用大量显存，降低推理效率。
• GPU预热：提前跑1-2次dummy推理，让CUDA kernel加载完成，消除首次推理的额外延迟。
• persistent_workers：开启后DataLoader的worker进程不会在每个epoch后销毁，提升后续批次的数据加载速度。

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