PyTorch模型训练效率监控方法

PyTorch模型训练效率监控方法

在实际项目中，我们经常遇到模型训练效率低下却无从下手的情况。本文通过具体案例分享几个实用的监控方法。

1. 使用torch.profiler进行性能分析

import torch
import torch.profiler

model = MyModel()
model.train()

with profiler.profile(
    activities=[profiler.ProfilerActivity.CPU, profiler.ProfilerActivity.CUDA],
    record_shapes=True
) as prof:
    for i, (inputs, targets) in enumerate(dataloader):
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

print(prof.key_averages().table(sort_by="self_cpu_time_total", row_limit=10))

2. 自定义训练循环监控

import time

class TrainingMonitor:
    def __init__(self):
        self.batch_times = []
        self.losses = []

    def record_batch(self, start_time, loss):
        batch_time = time.time() - start_time
        self.batch_times.append(batch_time)
        self.losses.append(loss)

# 使用示例
monitor = TrainingMonitor()
for epoch in range(10):
    for i, (inputs, targets) in enumerate(dataloader):
        start_time = time.time()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        
        monitor.record_batch(start_time, loss.item())

3. 实际测试数据

在1000个批次的训练中，发现以下问题：

• GPU利用率仅65%，主要瓶颈在数据加载
• 单批平均耗时：245ms (CPU: 180ms, GPU: 65ms)
• 内存峰值：1.2GB

通过优化数据加载器的num_workers=8后，效率提升约30%。

踩坑提醒：profiler会显著拖慢训练速度，建议仅在调试阶段使用。