使用torch.profiler发现DataLoader枚举耗时超50%，调参无效求帮助

问题背景 #

使用torch.profiler对fit_one_epoch函数性能分析时，发现DataLoader枚举阶段耗时占比超50%。尝试设置pin_memory=True或num_workers>0后，耗时占比未得到有效降低，寻求优化方案。

性能分析截图：

当前实现代码：

class NNModelTrainer:
    def __init__(self, model, optimizer, loss_func, n_epoch, batch_size, train_dataset, test_dataset):
        DEVICE = torch.device('cuda')
        self.model = model
        self.optimizer = optimizer
        self.loss_func = loss_func
        self.remaining_epochs = n_epoch
    
        self.train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
        self.test_dataset = test_dataset
    
    def fit_one_epoch(self):
        train_loss_val = 0
        self.model = self.model.to(DEVICE)
        self.model.train(mode=True)
            
        for x, y in self.train_dataloader:
            x, y = x.to(DEVICE), y.to(DEVICE)
            self.optimizer.zero_grad(set_to_none=False)
            yhat = self.model(x)
            train_loss = self.loss_func(yhat, y.squeeze(1)) * len(yhat)
            train_loss_val += train_loss.item()
            train_loss.backward()
            self.optimizer.step()
    
        self.remaining_epochs -= 1

优化方案 #

1. 修复变量作用域与冗余操作 #

代码中DEVICE是__init__局部变量，fit_one_epoch直接引用存在逻辑风险；且每次epoch重复执行model.to(DEVICE)属于冗余操作，建议调整为：

class NNModelTrainer:
    def __init__(self, model, optimizer, loss_func, n_epoch, batch_size, train_dataset, test_dataset):
        self.DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model = model.to(self.DEVICE)  # 仅初始化时执行一次
        self.optimizer = optimizer
        self.loss_func = loss_func
        self.remaining_epochs = n_epoch
    
        # DataLoader初始化后续调整
        self.train_dataloader = torch.utils.data.DataLoader(...)
    
    def fit_one_epoch(self):
        train_loss_val = 0
        self.model.train(mode=True)
            
        for x, y in self.train_dataloader:
            x, y = x.to(self.DEVICE), y.to(self.DEVICE)
            # 其余逻辑保持不变

2. 优化DataLoader核心参数配置 #

仅设置num_workers和pin_memory不足以解决问题，需配合以下参数：

• 匹配CPU核心数的num_workers: 建议设为os.cpu_count()或其1-2倍，避免过多进程导致调度开销
• 开启persistent_workers=True: 让worker进程在epoch间保持活跃，避免重复创建进程的开销
• pin_memory=True配合非阻塞传输: 数据移至GPU时添加non_blocking=True，进一步提升传输效率

调整后的DataLoader初始化代码：

import os

self.train_dataloader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=batch_size,
    shuffle=True,
    num_workers=os.cpu_count(),
    pin_memory=True,
    persistent_workers=True
)

# fit_one_epoch中数据传输修改为
x, y = x.to(self.DEVICE, non_blocking=True), y.to(self.DEVICE, non_blocking=True)

3. 优化Dataset预处理逻辑 #

DataLoader耗时过高往往源于__getitem__中CPU预处理开销过大，可从以下方向优化：

• 提前完成预处理：将数据预处理为张量格式保存，避免每次加载时重复计算
• 迁移部分预处理至GPU：使用torchvision.transforms.v2中支持GPU加速的变换操作
• 减少对象创建：在__getitem__中复用内存对象，降低频繁创建销毁的开销

4. 其他细节优化 #

• 调整预取因子：设置prefetch_factor=4（默认2），让worker提前预取更多批次数据，隐藏加载耗时
• 排查磁盘IO瓶颈：若为磁盘读取慢导致，可更换SSD存储或将常用数据集加载至内存缓存
• 按需关闭shuffle：若数据集已提前打乱，后续epoch可关闭shuffle减少采样开销

内容的提问来源于stack exchange，提问作者Ren