Python中PyTorch的Transform为什么导致CPU满载_使用Albumentations加速数据增强

PyTorch transforms.Compose CPU吃满主因是PIL操作单线程且每个DataLoader worker重复解码/转换；改用Albumentations（需转numpy、手动to_tensor、关is_check_shapes等）可显著降载。

PyTorch transforms.Compose 为什么会让 CPU 吃满？

根本原因不是 transform 本身慢，而是默认用 torchvision.transforms 做增强时，所有操作都在主线程（或 DataLoader 的 worker 进程）里同步执行 PIL 图像操作，且每个 worker 都要重复加载、解码、转换——尤其当用了 RandomRotation、ColorJitter 等基于 PIL 的变换时，CPU 解码 + 多通道 numpy 转换 + 再转回 tensor，开销极大。

更关键的是：PIL 操作是单线程的，无法利用多核；而 DataLoader(num_workers=N) 的每个 worker 又各自做一遍完整流程，导致 N 个 CPU 核心全在干重复的低效活。

常见现象：htop 显示 Python 进程 CPU 占用持续 90%+，GPU 利用率却只有 30%～50%，GPU 等数据等得发慌。

Albumentations 替换 torchvision.transforms 的实操要点

Albumentations 是专为 CV 数据增强设计的库，底层用 numba/cython 加速，所有操作原生支持 numpy array（无需 PIL 中转），且能 batch 化处理（虽然 DataLoader 本身不 batch，但单次调用效率高得多）。

• 必须把图像从 PIL → np.array（cv2.imread 或 np.array(pil_img)），Albumentations 不接受 PIL 对象
• 预处理 pipeline 要用 albumentations.Compose，不是 torchvision.transforms.Compose
• 输出仍是 np.ndarray，需手动转 torch.Tensor 并归一化（torch.from_numpy() + permute(2,0,1)）
• 别漏掉 to_tensor —— Albumentations 没有内置 tensor 转换，不像 torchvision 自带 ToTensorV2

示例片段：

Python 3.14.3

微软官方的 Python 扩展，是 VS Code 安装量最高的扩展（209M+）。集成 IntelliSense（通过 Pylance）、调试（通过 Python Debugger）、代码检查、格式化、重构和单元测试等功能。支持 Jupyter Notebook、虚拟环境管理和多 Python 版本切换。

import albumentations as A
import cv2
import torch
transform = A.Compose([
A.HorizontalFlip(p=0.5),
A.RandomBrightnessContrast(p=0.2),
A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
def load_and_augment(path):
img = cv2.imread(path)  # 直接读 BGR numpy array
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
augmented = transform(image=img)
tensor_img = torch.from_numpy(augmented["image"]).permute(2, 0, 1)
return tensor_img

DataLoader 配合 Albumentations 的关键配置

光换库不够，worker 行为必须调优。否则还是卡在 IO 和重复解码上。

• num_workers 别盲目设高——Albumentations 单 worker 吞吐更高，通常设 4 或 6 就够，再高反而因进程调度开销拖慢
• 务必加 pin_memory=True，配合 non_blocking=True 在 model.to(device) 时减少内存拷贝阻塞
• prefetch_factor 设为 2（默认是 2，但显式写出来更稳），避免 GPU 等 batch
• 禁用 torchvision.transforms.ToTensor 和 Normalize，它们和 Albumentations 的 Normalize 冲突，会 double 归一化

错误示范：transforms.Compose([ToTensor(), Normalize(...), albumentations_wrapper]) —— 这会导致像素值被缩放两次，模型训不动。

CPU 满载没缓解？检查这几个隐藏坑

换了 Albumentations 还卡，大概率是下面某个环节没断干净：

• 数据路径是 NFS 或远程挂载盘：IO 成瓶颈，cv2.imread 会卡住 worker，换成本地 SSD 或预加载到内存（小数据集可行）
• 用了 A.RandomCrop 但输入图尺寸远大于目标尺寸：Albumentations 内部仍要 alloc 大 buffer，内存带宽打满也会表现为 CPU 高占用
• 自定义 Dataset 的 __getitem__ 里写了 print/log/时间戳等调试代码：这些 I/O 操作在 worker 进程里串行执行，极易拖垮吞吐
• 没关掉 OpenCV 的多线程：OpenCV 默认启用 TBB，和 DataLoader worker 冲突，加 cv2.setNumThreads(0) 到初始化位置

最易忽略的一点：Albumentations 的 Compose 默认开启 is_check_shapes=True，每次都会校验 image/mask 的 shape 对齐——对纯分类任务完全没必要，设 is_check_shapes=False 能省 5%～10% CPU 时间。