Ubuntu下PyTorch內存怎么優化

Ubuntu下PyTorch內存優化方法

1. 降低批次大?。˙atch Size）

批次大小是影響GPU顯存使用的核心因素之一。較小的批次大小能直接減少顯存占用，但需平衡其對訓練速度（如梯度更新頻率）和模型性能（如泛化能力）的影響。建議通過實驗找到“顯存占用可接受且不影響模型效果”的最小批次值。

2. 使用半精度浮點數（Float16）與自動混合精度（AMP）

半精度（float16）相比單精度（float32）可減少50%的顯存占用，同時通過PyTorch的torch.cuda.amp模塊實現自動混合精度（AMP），能在保持模型數值穩定性的前提下，自動在float16和float32之間切換（如梯度計算用float32保證穩定性，前向/反向傳播用float16提升速度）。示例代碼：

from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
for inputs, targets in dataloader:
    optimizer.zero_grad()
    with autocast():  # 自動混合精度上下文
        outputs = model(inputs)
        loss = criterion(outputs, targets)
    scaler.scale(loss).backward()  # 縮放梯度避免underflow
    scaler.step(optimizer)         # 更新參數
    scaler.update()                # 調整縮放因子

3. 釋放不必要的張量與緩存

• 手動刪除無用張量：使用del關鍵字刪除不再需要的中間變量（如損失值、預測結果），斷開其對顯存的引用。
• 清空GPU緩存：調用torch.cuda.empty_cache()釋放PyTorch緩存的無用顯存（如未使用的中間結果），注意該操作不會釋放被引用的張量。
• 觸發垃圾回收：使用gc.collect()強制Python垃圾回收器回收無用對象，配合del和empty_cache()效果更佳。
  示例代碼：

del outputs, loss  # 刪除無用變量
torch.cuda.empty_cache()  # 清空GPU緩存
import gc
gc.collect()  # 觸發垃圾回收

4. 梯度累積（Gradient Accumulation）

梯度累積通過“多次小批次計算梯度→累加→一次更新”的方式，模擬更大批次的效果，同時不增加顯存占用。適用于“顯存不足但需較大批次”的場景。示例代碼：

accumulation_steps = 4  # 累積4個小批次的梯度
for i, (inputs, targets) in enumerate(dataloader):
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    loss = loss / accumulation_steps  # 歸一化損失（避免梯度爆炸）
    loss.backward()  # 累積梯度
    
    if (i + 1) % accumulation_steps == 0:  # 達到累積步數時更新參數
        optimizer.step()
        optimizer.zero_grad()  # 清零梯度

5. 優化數據加載流程

• 增加數據加載并行性：通過DataLoader的num_workers參數（設置為CPU核心數的2-4倍）啟用多進程數據加載，避免數據預處理成為瓶頸。
• 減少內存拷貝：確保數據預處理后的結果（如numpy數組轉tensor）直接存入GPU，避免CPU-GPU間的多次數據傳輸。
• 使用高效存儲格式：對于大型數據集，采用HDF5或LMDB等格式，減少內存占用和加載時間。

6. 使用內存高效的模型結構

• 替換全連接層：用卷積層代替全連接層（如將FC層的in_features與out_features轉換為卷積核的in_channels與out_channels），減少參數數量（如ResNet-50的全連接層參數占比約90%）。
• 采用深度可分離卷積：將標準卷積分解為“深度卷積（Depthwise Convolution）+ 逐點卷積（Pointwise Convolution）”，降低計算量和顯存占用（如MobileNet系列模型）。

7. 分布式訓練（Distributed Training）

將模型訓練分布到多個GPU（單機多卡）或多臺機器（多機多卡），通過數據并行（Data Parallelism）或模型并行（Model Parallelism）減少單個設備的顯存負載。PyTorch提供torch.distributed模塊支持分布式訓練，常用 launch 工具如torchrun。示例代碼（數據并行）：

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group(backend='nccl')  # 初始化進程組
model = DDP(model.to(device))  # 包裝模型

8. 監控與診斷內存使用

• 查看顯存占用：使用torch.cuda.memory_summary()打印顯存分配詳情（如已用/剩余顯存、緩存情況），或nvidia-smi命令實時監控GPU顯存使用率。
• 定位內存泄漏：通過torch.utils.checkpoint檢查張量是否意外保留計算圖（如非訓練場景未用with torch.no_grad()），或使用memory_profiler庫逐行跟蹤內存變化。

9. 系統級優化

• 清理系統緩存：定期執行sudo echo 3 | sudo tee /proc/sys/vm/drop_caches釋放系統頁面緩存（不影響正在運行的程序）。
• 關閉后臺程序：停止不必要的系統服務（如桌面特效、自動更新），釋放系統內存。
• 設置虛擬內存（Swap）：若物理內存不足，創建Swap文件（如sudo dd if=/dev/zero of=/swapfile bs=64M count=16，sudo mkswap /swapfile，sudo swapon /swapfile），作為物理內存的擴展（注意：Swap性能低于物理內存，僅作臨時解決方案）。