[分享] PyTorch 訓練目前已經是 num_workers = 0 了 CPU占用還是高

PyTorch 訓練目前已經是 num_workers = 0 了 CPU占用還是高

如果 num_workers=0 後 CPU 占用還是很高,問題主要在於模型運算本身。以下是幾個有效的解決方案:

1. 限制 PyTorch 的 CPU 執行緒數 (最有效)

在 training.py 最前面加入:

import torch
import os

# === 在所有 import 之後,主程式之前加入 ===
# 限制 PyTorch 使用的 CPU 執行緒數
torch.set_num_threads(2) # 限制為 2 個執行緒
torch.set_num_interop_threads(2) # 限制 inter-op 執行緒

# 限制底層數學庫的執行緒數
os.environ['OMP_NUM_THREADS'] = '2'
os.environ['MKL_NUM_THREADS'] = '2'
os.environ['NUMEXPR_NUM_THREADS'] = '2'
os.environ['OPENBLAS_NUM_THREADS'] = '2'

2. 在 CONFIG 中加入可配置的執行緒控制

修改 training.py 的 main 函數開頭:

def main():
  # === 在最開始加入執行緒控制 ===
  cpu_threads = CONFIG.get('cpu_threads', 2)
  torch.set_num_threads(cpu_threads)
  torch.set_num_interop_threads(cpu_threads)
  os.environ['OMP_NUM_THREADS'] = str(cpu_threads)
  os.environ['MKL_NUM_THREADS'] = str(cpu_threads)
  os.environ['NUMEXPR_NUM_THREADS'] = str(cpu_threads)
  os.environ['OPENBLAS_NUM_THREADS'] = str(cpu_threads)
  
  print(f"? 限制 CPU 使用 {cpu_threads} 個執行緒")
  
  print("="*70)
  print("影像分類訓練系統")
  # ... 其餘程式碼

在 CONFIG 設定中加入:

CONFIG = {
  # ... 其他設定 ...
  'cpu_threads': config.get('cpu_threads', 2), # 新增這行
}

3. 在 settings/training.yaml 加入設定

# CPU 執行緒控制
cpu_threads: 2 # 可調整為 1, 2, 4 等
num_workers: 0

4. 減小 batch_size (如果記憶體夠用)

較小的 batch_size 會降低每次運算的負擔:

batch_size: 16 # 從 32 改為 16 (或更小)

5. 加入訓練間隔休息 (可選)

在 train_epoch 函數中加入小延遲:

def train_epoch(model, dataloader, criterion, optimizer, device):
  model.train()
  running_loss = 0.0
  total_batches = len(dataloader)
  
  for batch_idx, (inputs, labels) in enumerate(dataloader, 1):
    inputs, labels = inputs.to(device), labels.float().to(device)
    
    optimizer.zero_grad()
    outputs = model(inputs).squeeze(1)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()
    
    running_loss += loss.item() * inputs.size(0)
    
    # 計算當前平均 loss
    current_avg_loss = running_loss / (batch_idx * inputs.size(0))
    
    # 顯示訓練進度
    progress = batch_idx / total_batches * 100
    print(f"\r Training: [{batch_idx}/{total_batches}] {progress:.1f}% | Loss: {current_avg_loss:.4f}", end='')
    
    # === 每 N 個 batch 休息一下 (可選) ===
    if batch_idx % 10 == 0:
      import time
      time.sleep(0.01) # 休息 10ms
  
  epoch_loss = running_loss / len(dataloader.dataset)
  print()
  return epoch_loss

6. 使用更輕量的模型

目前使用 efficientnet_v2_s,可以改為更小的模型:

model: mobilenet_v3_small # 參數量只有 2.5M,比 efficientnet_v2_s (21M) 小很多

完整修改建議:

在 if __name__ == '__main__': 之前加入:

if __name__ == '__main__':
  import torch
  import os
  
  # === 必須在最開始設定,否則無效 ===
  torch.set_num_threads(2)
  torch.set_num_interop_threads(2)
  os.environ['OMP_NUM_THREADS'] = '2'
  os.environ['MKL_NUM_THREADS'] = '2'
  os.environ['NUMEXPR_NUM_THREADS'] = '2'
  os.environ['OPENBLAS_NUM_THREADS'] = '2'

  warnings.filterwarnings('ignore')
  # 載入設定檔
  config_path = 'settings/training.yaml'
  # ... 其餘程式碼

在 settings/training.yaml:

# 效能調整
cpu_threads: 2   # CPU 執行緒數 (1-4)
num_workers: 0   # DataLoader 工作執行緒
batch_size: 16   # 減小 batch size
model: mobilenet_v3_small # 使用更輕量的模型

效果預期:

torch.set_num_threads(2): 可降低 50-70% 的 CPU 占用

改用 mobilenet_v3_small: 可再降低 20-30% 的運算量

減小 batch_size: 可降低瞬間峰值,但訓練時間會拉長

建議優先嘗試第 1 點,這是最有效且最簡單的方法!

後記:

# 限制 PyTorch 使用的 CPU 執行緒數
torch.set_num_threads(2)
torch.set_num_interop_threads(2)
os.environ['OMP_NUM_THREADS'] = '2'
os.environ['MKL_NUM_THREADS'] = '2'
os.environ['NUMEXPR_NUM_THREADS'] = '2'
os.environ['OPENBLAS_NUM_THREADS'] = '2'

實測有效!

Zoearth Joomla Site
台中阿任的Joomla網站