CPU usage is too high, and the GPU cannot be utilized
Hello, I noticed during training that adjusting the batch size does not affect GPU memory usage, and my CPU usage is very high. Could you please explain why this is?
期待作者的回应┭┮﹏┭┮
我也遇到了同一问题。异常检测任务在多进程并行下是cpu依赖型的,而且作者的代码实现对cpu不友好。你可手动优化
我用claude分析了下:
## PromptAD CPU 瓶颈分析报告
根据代码分析,发现了 **5 个主要的 CPU 密集型瓶颈**:
### 🔴 **1. DataLoader 设置 `num_workers=0`(最严重)**
**位置**: __init__.py 第 38、41 行
```python
data_loader = DataLoader(dataset_inst, batch_size=kwargs['batch_size'],
shuffle=True, num_workers=0) # ❌ 单线程加载
问题:
- 所有数据加载都在主进程中完成,完全串行
- GPU 在等待 CPU 准备下一批数据时处于空闲状态
- 无法利用多核 CPU 并行预处理数据
影响: 最严重的瓶颈,导致 GPU 利用率低下
🔴 2. 训练循环中实时图像转换(次严重)
位置: train_cls.py 第 45、63 行,train_seg.py 第 44、64 行
# 每个 epoch 每个 batch 都要执行:
for (data, mask, label, name, img_type) in train_data:
# ❌ CPU 密集型操作在训练循环中
data = [model.transform(Image.fromarray(cv2.cvtColor(f.numpy(), cv2.COLOR_BGR2RGB)))
for f in data]
data = torch.stack(data, dim=0).to(device)
问题分析:
f.numpy(): Tensor → NumPy 转换(CPU)cv2.cvtColor(BGR→RGB): 颜色空间转换(CPU)Image.fromarray(): NumPy → PIL Image(CPU)model.transform(): CLIP 预处理(Resize, CenterCrop, Normalize)(CPU)torch.stack(): 重新组合 Tensor(CPU)
影响:
- 每个样本都要经历 5 次 CPU 操作
- batch_size=8, epoch=100 → 每个类别 800 次循环
- k-shot 训练集只有几张图,却每个 epoch 都重新转换
🔴 3. Dataset 中的高分辨率 resize(中等)
位置: dataset.py 第 46-47 行
def __getitem__(self, idx):
img = cv2.imread(img_path, cv2.IMREAD_COLOR)
img = cv2.resize(img, (1024, 1024)) # ❌ 每次读取都 resize 到 1024
gt = cv2.resize(gt, (1024, 1024), interpolation=cv2.INTER_NEAREST)
问题:
- MVTec 原始图像约 700×700,VisA 约 900×1200
- 先 resize 到 1024×1024,然后在
model.transform()中再 resize 到 224×224 - 双重 resize,且中间分辨率过高
影响:
- 浪费 CPU 计算 1024×1024 的数据,最终只用 224×224
cv2.resize是 CPU 密集型操作
🔴 4. 缺少 pin_memory 优化(轻微)
位置: init.py DataLoader 配置
DataLoader(..., num_workers=0) # 缺少 pin_memory=True
问题:
- 数据从 CPU 转移到 GPU 时需要经过分页内存
- 没有使用固定内存(pinned memory)加速传输
🔴 5. 评估阶段的重复转换(轻微)
位置: train_cls.py 第 126 行
for (data, mask, label, name, img_type) in dataloader:
data = [model.transform(Image.fromarray(f.numpy())) for f in data] # ❌ 再次转换
问题: 测试集每个 epoch 都重新转换一次
📊 为什么 CPU 成为瓶颈?
从你的 htop 截图分析:
- 40 个核心全部处于 ~60-90% 使用率 → CPU 非常忙碌
- 每个任务独占 4-6 个核心 → 但
num_workers=0实际只用 1 个核心 - 3 个 GPU 利用率很低(18-20%) → GPU 在等待数据
实际情况:
训练流程:
[CPU] 读取图像 → resize 1024 →
[CPU] 等待 →
[CPU] cvtColor BGR→RGB →
[CPU] 转 PIL Image →
[CPU] transform (resize 224, normalize) →
[CPU] stack tensors →
[GPU] 等待... →
[GPU] 模型前向传播 (很快) →
[GPU] 空闲等待下一批数据...
因为 num_workers=0,GPU 大部分时间在等待 CPU 的单线程数据预处理。
✅ 优化建议(按优先级排序)
🥇 优先级 1: 启用多进程数据加载
# datasets/__init__.py
DataLoader(dataset_inst,
batch_size=kwargs['batch_size'],
shuffle=True,
num_workers=4, # ⭐ 使用 4 个子进程
pin_memory=True, # ⭐ 加速 GPU 传输
persistent_workers=True) # ⭐ 保持 worker 进程活跃
预期提升: GPU 利用率 20% → 70%+
🥈 优先级 2: 将图像预处理移到 Dataset
# datasets/dataset.py
class CLIPDataset(Dataset):
def __init__(self, ..., transform=None):
self.transform = transform # 传入 CLIP transform
def __getitem__(self, idx):
img = cv2.imread(img_path, cv2.IMREAD_COLOR)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 在这里转换
img = Image.fromarray(img)
if self.transform:
img = self.transform(img) # ⭐ Dataset 中完成 transform
# 返回处理好的 Tensor,而非原始数据
return img, gt, label, img_name, img_type
然后在 train_cls.py 中移除重复转换:
# ❌ 删除这些:
# data = [model.transform(Image.fromarray(cv2.cvtColor(f.numpy(), ...)))]
# ✅ 直接使用:
for (data, mask, label, name, img_type) in train_data:
data = data.to(device) # 已经是 Tensor
预期提升: 减少 80% 的 CPU 开销
🥉 优先级 3: 降低中间 resize 分辨率
# datasets/dataset.py
img = cv2.resize(img, (256, 256)) # 从 1024 降到 256
# CLIP transform 会进一步 resize 到 224
预期提升: 减少 60% 的 resize 开销
💡 优先级 4: 使用数据预缓存(可选)
对于 k-shot(只有几张训练图像),可以在 __init__ 中预加载所有数据到内存:
class CLIPDataset(Dataset):
def __init__(self, ..., preload=True):
if preload and phase == 'train':
self.cached_data = [self._load_and_transform(i) for i in range(len(self))]
📈 预期整体提升
应用上述优化后:
- GPU 利用率: 20% → 80%+
- 训练速度: 提升 3-5 倍
- CPU 占用: 下降至合理水平
- 并行能力: 可以运行更多实验(目前 CPU 是限制因素)
你想让我实现这些优化吗?我可以创建优化版本的代码。
经验表明num_workers设置为1时表现最好。应用以上修改后,单任务cpu的负载下降约50-60%