【Bug已解决】Pinning memory is actually slower in PyTorch? 解决方案问题描述在 PyTorch 中pin_memory页锁定内存是一个被广泛推荐的数据加载优化技术。官方文档和大量教程都建议在DataLoader中设置pin_memoryTrue以加速从 CPU 到 GPU 的数据传输。然而许多开发者发现在实际使用中开启pin_memory后训练速度不仅没有提升反而变慢了。这种现象看似矛盾——一个本应加速训练的优化手段却导致了性能退化。问题的根源涉及多个层面操作系统内存管理机制、PyTorch 的内存分配策略、数据加载器的工作方式以及硬件配置等。本文将深入分析为什么pin_memory在某些情况下会变慢并提供完整的诊断和解决方案。错误复现错误场景一pin_memory 导致训练变慢import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader import time import numpy as np # 创建一个简单的数据集 class SimpleDataset(Dataset): def __init__(self, num_samples10000, data_size224*224*3): self.data np.random.randn(num_samples, data_size).astype(np.float32) self.labels np.random.randint(0, 10, num_samples) def __len__(self): return len(self.labels) def __getitem__(self, idx): return torch.from_numpy(self.data[idx]), self.labels[idx] dataset SimpleDataset(num_samples10000) # 测试 pin_memoryFalse dataloader_no_pin DataLoader(dataset, batch_size64, shuffleTrue, num_workers4, pin_memoryFalse) # 测试 pin_memoryTrue dataloader_pin DataLoader(dataset, batch_size64, shuffleTrue, num_workers4, pin_memoryTrue) # 简单模型 model nn.Sequential( nn.Linear(224*224*3, 512), nn.ReLU(), nn.Linear(512, 10) ).cuda() criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.parameters()) def train_epoch(dataloader, model, criterion, optimizer): model.train() total_time 0 for batch_x, batch_y in dataloader: start time.time() batch_x batch_x.cuda(non_blockingTrue) batch_y batch_y.cuda(non_blockingTrue) optimizer.zero_grad() output model(batch_x) loss criterion(output, batch_y) loss.backward() optimizer.step() total_time time.time() - start return total_time # 对比测试 print(Testing without pin_memory...) time_no_pin train_epoch(dataloader_no_pin, model, criterion, optimizer) print(fTime without pin_memory: {time_no_pin:.2f}s) print(\nTesting with pin_memory...) time_pin train_epoch(dataloader_pin, model, criterion, optimizer) print(fTime with pin_memory: {time_pin:.2f}s) print(f\nDifference: pin_memory is {slower if time_pin time_no_pin else faster} by {abs(time_pin - time_no_pin):.2f}s)运行结果可能如下Testing without pin_memory... Time without pin_memory: 12.34s Testing with pin_memory... Time with pin_memory: 15.67s Difference: pin_memory is slower by 3.33s错误场景二内存不足导致频繁交换# 当数据集很大时pin_memory 会占用大量页锁定内存 # 导致操作系统频繁进行内存交换swap严重拖慢速度 large_dataset SimpleDataset(num_samples100000, data_size512*512*3) # pin_memoryTrue 会尝试将所有数据复制到页锁定内存 # 如果系统物理内存不足会导致严重的性能问题 dataloader DataLoader(large_dataset, batch_size32, shuffleTrue, num_workers8, pin_memoryTrue) # 系统可能报出内存警告或 OOM try: for batch_x, batch_y in dataloader: batch_x batch_x.cuda() # ... 训练代码 except RuntimeError as e: print(fError: {e}) # 可能的输出: CUDA out of memory 或 系统级别的 OOM错误场景三num_workers 与 pin_memory 的冲突# 在某些系统上多进程 pin_memory 会导致竞争和锁争用 dataloader DataLoader( dataset, batch_size64, shuffleTrue, num_workers16, # 过多的 worker pin_memoryTrue, persistent_workersTrue ) # 在 macOS 或某些 Linux 配置上这可能导致 # 1. worker 进程之间的内存竞争 # 2. 页锁定内存分配失败 # 3. 数据加载变慢而非变快根因分析1. 什么是页锁定内存Pinned Memory在操作系统中内存分为两种可分页内存Pageable Memory普通内存操作系统可以将其换出到磁盘swap space以释放物理内存。页锁定内存Pinned/Page-locked Memory被锁定在物理内存中不会被操作系统换出。GPU 通过 DMADirect Memory Access技术直接从 CPU 内存读取数据。DMA 操作要求源内存是页锁定的因为 DMA 操作不能触发缺页中断。当你将可分页内存传输到 GPU 时驱动程序需要先分配一块页锁定内存将数据从可分页内存复制到页锁定内存通过 DMA 将数据从页锁定内存传输到 GPU释放页锁定内存而使用页锁定内存时步骤 1、2、4 被省略理论上应该更快。2. 为什么 pin_memory 反而变慢原因一内存分配开销pin_memoryTrue时DataLoader 会在主进程中创建一个pin_memory_thread该线程负责将 worker 进程返回的可分页内存数据复制到页锁定内存。这个复制操作本身是有开销的# DataLoader 内部大致做了以下操作 # 1. worker 进程返回可分页内存的 batch # 2. pin_memory_thread 分配页锁定内存 # 3. pin_memory_thread 将数据复制到页锁定内存 # 4. 返回页锁定内存的 batch如果数据量较小或模型计算量较小这个额外的复制开销可能超过 DMA 传输节省的时间。原因二页锁定内存是稀缺资源操作系统的页锁定内存总量是有限的。当pin_memory占用过多页锁定内存时其他系统进程无法使用这些内存操作系统可能被迫进行更多的内存管理操作在内存紧张时系统性能急剧下降原因三与non_blocking的误用pin_memory的真正优势在于配合non_blockingTrue使用时可以实现数据传输和计算的重叠# 正确用法pin_memory non_blocking 实现重叠 batch_x batch_x.cuda(non_blockingTrue) # 异步传输 # 此时 CPU 可以继续执行其他操作GPU 在后台传输数据但如果你的代码中没有使用non_blockingTrue或者数据传输完成后立即需要使用数据没有可重叠的计算pin_memory的优势就无法体现。原因四CPU 瓶颈当训练的瓶颈在 CPU 端数据预处理、数据增强等而非 GPU 端时pin_memory的额外 CPU 开销会加剧瓶颈导致整体变慢。原因五小 batch size 或小数据对于小 batch size 或小尺寸数据DMA 传输本身非常快微秒级页锁定内存的分配和复制开销可能超过传输节省的时间。原因六多 worker 竞争当num_workers较多时多个 worker 进程同时返回数据pin_memory_thread需要串行处理这些数据可能成为瓶颈。同时多个进程竞争页锁定内存分配也可能导致性能下降。3. 性能影响的量化分析因素pin_memory 有利pin_memory 不利Batch Size大 batch128小 batch32数据大小大数据图像、视频小数据文本、表格GPU 利用率GPU 计算密集GPU 空闲等待系统内存充足32GB不足16GBnum_workers少1-4多8non_blocking使用不使用解决方案方案一正确使用 pin_memory non_blockingimport torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader import time class ImageDataset(Dataset): def __init__(self, num_samples5000, size3*224*224): self.data torch.randn(num_samples, size) self.labels torch.randint(0, 10, (num_samples,)) def __len__(self): return len(self.labels) def __getitem__(self, idx): return self.data[idx], self.labels[idx] class SimpleModel(nn.Module): def __init__(self): super().__init__() self.fc nn.Sequential( nn.Linear(3*224*224, 1024), nn.ReLU(), nn.Linear(1024, 512), nn.ReLU(), nn.Linear(512, 10) ) def forward(self, x): return self.fc(x) def benchmark(dataloader, model, use_non_blocking, num_epochs3): model.train() criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.parameters()) total_time 0 for epoch in range(num_epochs): epoch_start time.time() for batch_x, batch_y in dataloader: # 关键使用 non_blockingTrue 实现异步传输 batch_x batch_x.cuda(non_blockinguse_non_blocking) batch_y batch_y.cuda(non_blockinguse_non_blocking) optimizer.zero_grad() output model(batch_x) loss criterion(output, batch_y) loss.backward() optimizer.step() total_time time.time() - epoch_start return total_time / num_epochs # 实验设置 dataset ImageDataset(num_samples5000) model SimpleModel().cuda() configs [ (No pin, no non_blocking, False, False, False), (pin_memory, no non_blocking, True, False, True), (pin_memory non_blocking, True, True, True), ] print(Benchmarking different configurations:\n) for name, pin_mem, non_block, _ in configs: dataloader DataLoader( dataset, batch_size128, shuffleTrue, num_workers4, pin_memorypin_mem ) avg_time benchmark(dataloader, model, non_block) print(f{name}: {avg_time:.2f}s per epoch)方案二根据场景选择是否启用 pin_memoryimport torch from torch.utils.data import DataLoader def create_optimal_dataloader(dataset, batch_size, model_typeimage, data_size_bytesNone, system_memory_gb16): 根据场景自动选择最优的 DataLoader 配置。 Args: dataset: 数据集 batch_size: 批大小 model_type: image图像, text文本, tabular表格 data_size_bytes: 单个样本的字节数 system_memory_gb: 系统物理内存GB Returns: 配置好的 DataLoader # 估算每 batch 的内存占用 if data_size_bytes: batch_memory batch_size * data_size_bytes else: if model_type image: batch_memory batch_size * 3 * 224 * 224 * 4 # float32 elif model_type text: batch_memory batch_size * 512 * 4 else: batch_memory batch_size * 100 * 4 # 决定是否使用 pin_memory # 条件batch 较大、数据较大、系统内存充足 use_pin_memory ( batch_size 64 and model_type image and system_memory_gb 16 and batch_memory system_memory_gb * 1024**3 * 0.3 # 不超过系统内存的30% ) # 决定 num_workers cpu_count torch.get_num_threads() if model_type image: num_workers min(4, cpu_count) else: num_workers min(2, cpu_count) print(fConfiguration:) print(f pin_memory: {use_pin_memory}) print(f num_workers: {num_workers}) print(f Estimated batch memory: {batch_memory / 1024 / 1024:.2f} MB) return DataLoader( dataset, batch_sizebatch_size, shuffleTrue, num_workersnum_workers, pin_memoryuse_pin_memory, persistent_workersnum_workers 0, prefetch_factor2 if num_workers 0 else None ) # 使用示例 # dataloader create_optimal_dataloader(dataset, batch_size128, model_typeimage)方案三手动管理页锁定内存对于需要精细控制的高级用户import torch def manual_pin_transfer(tensor, devicecuda, non_blockingTrue): 手动管理页锁定内存传输。 适用于需要精细控制内存使用的场景。 if not tensor.is_pinned(): # 分配页锁定内存并复制 pinned_tensor tensor.pin_memory() else: pinned_tensor tensor # 异步传输到 GPU return pinned_tensor.to(device, non_blockingnon_blocking) # 使用示例 data torch.randn(64, 3, 224, 224) # 方式1直接传输不使用 pin_memory # data_gpu data.cuda() # 方式2手动 pin 后传输 data_gpu manual_pin_transfer(data, devicecuda, non_blockingTrue) # 检查是否为页锁定内存 print(fIs pinned: {data.is_pinned()}) # False pinned data.pin_memory() print(fIs pinned: {pinned.is_pinned()}) # True完整修复代码以下是一个完整的、可运行的训练脚本包含 pin_memory 性能诊断和优化import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import Dataset, DataLoader import time import numpy as np import os # # 1. 数据集定义 # class SyntheticImageDataset(Dataset): 合成图像数据集用于性能测试。 模拟真实图像分类场景的数据加载。 def __init__(self, num_samples10000, channels3, height224, width224, num_classes10, augmentFalse): self.num_samples num_samples self.channels channels self.height height self.width width self.num_classes num_classes self.augment augment # 不预加载所有数据模拟真实场景 self.labels np.random.randint(0, num_classes, num_samples) def __len__(self): return self.num_samples def __getitem__(self, idx): # 模拟数据加载和预处理 data np.random.randn(self.channels, self.height, self.width).astype(np.float32) if self.augment: # 模拟数据增强开销 data self._augment(data) label self.labels[idx] return torch.from_numpy(data), torch.tensor(label, dtypetorch.long) def _augment(self, data): 简单的数据增强 # 随机水平翻转 if np.random.random() 0.5: data data[:, :, ::-1].copy() # 随机裁剪 h_off np.random.randint(0, 10) w_off np.random.randint(0, 10) data data[:, h_off:h_offself.height-10, w_off:w_offself.width-10] # 重新调整大小简化版 padded np.zeros((self.channels, self.height, self.width), dtypenp.float32) padded[:, :data.shape[1], :data.shape[2]] data return padded # # 2. 模型定义 # class ResNetLikeModel(nn.Module): 简化的 ResNet 风格模型用于性能测试 def __init__(self, num_classes10): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 64, 3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(64, 128, 3, padding1), nn.BatchNorm2d(128), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(128, 256, 3, padding1), nn.BatchNorm2d(256), nn.ReLU(inplaceTrue), nn.AdaptiveAvgPool2d((1, 1)) ) self.classifier nn.Sequential( nn.Flatten(), nn.Linear(256, 128), nn.ReLU(), nn.Dropout(0.5), nn.Linear(128, num_classes) ) def forward(self, x): x self.features(x) x self.classifier(x) return x # # 3. 性能诊断工具 # class DataLoaderBenchmark: DataLoader 性能诊断工具。 帮助确定 pin_memory 是否有益。 def __init__(self, dataset, batch_size64, modelNone, devicecuda): self.dataset dataset self.batch_size batch_size self.device device self.model model or ResNetLikeModel().to(device) self.criterion nn.CrossEntropyLoss() self.optimizer optim.Adam(self.model.parameters()) def _train_epoch(self, dataloader, use_non_blockingTrue): self.model.train() data_times [] compute_times [] total_start time.time() for batch_x, batch_y in dataloader: data_start time.time() batch_x batch_x.to(self.device, non_blockinguse_non_blocking) batch_y batch_y.to(self.device, non_blockinguse_non_blocking) # 等待数据传输完成 if use_non_blocking: torch.cuda.synchronize() data_time time.time() - data_start data_times.append(data_time) compute_start time.time() self.optimizer.zero_grad() output self.model(batch_x) loss self.criterion(output, batch_y) loss.backward() self.optimizer.step() torch.cuda.synchronize() compute_time time.time() - compute_start compute_times.append(compute_time) total_time time.time() - total_start return { total: total_time, data_transfer_avg: np.mean(data_times), compute_avg: np.mean(compute_times), data_transfer_total: np.sum(data_times), compute_total: np.sum(compute_times), } def run_comparison(self, num_workers4, num_epochs3): 对比不同配置的性能。 configs [ {name: No pin_memory, pin_memory: False, non_blocking: False}, {name: pin_memory only, pin_memory: True, non_blocking: False}, {name: pin_memory non_blocking, pin_memory: True, non_blocking: True}, ] results {} for config in configs: print(f\n--- Testing: {config[name]} ---) dataloader DataLoader( self.dataset, batch_sizeself.batch_size, shuffleTrue, num_workersnum_workers, pin_memoryconfig[pin_memory], persistent_workersTrue ) # 预热 print(Warming up...) self._train_epoch(dataloader, config[non_blocking]) # 正式测试 print(Benchmarking...) epoch_results [] for epoch in range(num_epochs): result self._train_epoch(dataloader, config[non_blocking]) epoch_results.append(result) print(f Epoch {epoch1}: {result[total]:.2f}s f(data: {result[data_transfer_avg]*1000:.1f}ms/batch, fcompute: {result[compute_avg]*1000:.1f}ms/batch)) avg_total np.mean([r[total] for r in epoch_results]) avg_data np.mean([r[data_transfer_avg] for r in epoch_results]) avg_compute np.mean([r[compute_avg] for r in epoch_results]) results[config[name]] { avg_total: avg_total, avg_data_transfer: avg_data, avg_compute: avg_compute, } # 打印汇总 print(\n * 60) print(BENCHMARK SUMMARY) print( * 60) for name, result in results.items(): print(f{name:35s}: {result[avg_total]:.2f}s/epoch f(data: {result[avg_data_transfer]*1000:.1f}ms, fcompute: {result[avg_compute]*1000:.1f}ms)) # 推荐配置 fastest min(results.items(), keylambda x: x[1][avg_total]) print(f\nRecommended: {fastest[0]}) return results # # 4. 优化的训练函数 # def optimized_train(model, train_dataset, val_dataset, num_epochs10, devicecuda, batch_size128): 优化的训练函数自动选择最佳 pin_memory 配置。 # 先运行基准测试确定最佳配置 print(Running DataLoader benchmark to determine optimal configuration...) benchmark DataLoaderBenchmark( train_dataset, batch_sizebatch_size, modelmodel, devicedevice ) results benchmark.run_comparison(num_workers4, num_epochs2) # 选择最快的配置 fastest_config min(results.items(), keylambda x: x[1][avg_total]) use_pin_memory pin_memory in fastest_config[0].lower() use_non_blocking non_blocking in fastest_config[0].lower() print(f\nUsing configuration: pin_memory{use_pin_memory}, fnon_blocking{use_non_blocking}) # 创建最优的 DataLoader train_loader DataLoader( train_dataset, batch_sizebatch_size, shuffleTrue, num_workers4, pin_memoryuse_pin_memory, persistent_workersTrue, prefetch_factor2 ) val_loader DataLoader( val_dataset, batch_sizebatch_size, shuffleFalse, num_workers2, pin_memoryuse_pin_memory, persistent_workersTrue ) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.001) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_maxnum_epochs) best_val_acc 0.0 for epoch in range(num_epochs): # 训练 model.train() train_loss 0.0 train_correct 0 train_total 0 epoch_start time.time() for batch_x, batch_y in train_loader: batch_x batch_x.to(device, non_blockinguse_non_blocking) batch_y batch_y.to(device, non_blockinguse_non_blocking) optimizer.zero_grad() output model(batch_x) loss criterion(output, batch_y) loss.backward() optimizer.step() train_loss loss.item() * batch_x.size(0) _, predicted output.max(1) train_correct predicted.eq(batch_y).sum().item() train_total batch_y.size(0) train_loss / train_total train_acc train_correct / train_total epoch_time time.time() - epoch_start # 验证 model.eval() val_loss 0.0 val_correct 0 val_total 0 with torch.no_grad(): for batch_x, batch_y in val_loader: batch_x batch_x.to(device, non_blockinguse_non_blocking) batch_y batch_y.to(device, non_blockinguse_non_blocking) output model(batch_x) loss criterion(output, batch_y) val_loss loss.item() * batch_x.size(0) _, predicted output.max(1) val_correct predicted.eq(batch_y).sum().item() val_total batch_y.size(0) val_loss / val_total val_acc val_correct / val_total scheduler.step() if val_acc best_val_acc: best_val_acc val_acc torch.save(model.state_dict(), best_model.pth) print(fEpoch [{epoch1}/{num_epochs}] ({epoch_time:.1f}s) fTrain Loss: {train_loss:.4f}, Acc: {train_acc:.4f} | fVal Loss: {val_loss:.4f}, Acc: {val_acc:.4f}) return model # # 5. 主程序 # def main(): torch.manual_seed(42) np.random.seed(42) device torch.device(cuda if torch.cuda.is_available() else cpu) print(fDevice: {device}) print(fCPU threads: {torch.get_num_threads()}) if torch.cuda.is_available(): print(fGPU: {torch.cuda.get_device_name()}) print(fGPU Memory: {torch.cuda.get_device_properties(0).total_mem / 1024**3:.1f} GB) # 创建数据集 train_dataset SyntheticImageDataset(num_samples2000, augmentTrue) val_dataset SyntheticImageDataset(num_samples400, augmentFalse) # 创建模型 model ResNetLikeModel(num_classes10).to(device) # 运行优化训练 model optimized_train( model, train_dataset, val_dataset, num_epochs5, devicedevice, batch_size64 ) print(\nTraining complete!) if __name__ __main__: main()常见陷阱与注意事项1. pin_memory 不等于自动加速pin_memoryTrue只是一个可能加速的设置不是保证加速。它的效果取决于数据大小、batch size、系统内存、GPU 利用率等多个因素。盲目开启可能导致性能下降。2. 必须配合 non_blocking 使用pin_memory的真正优势在于配合non_blockingTrue实现数据传输和计算的重叠。如果只设置pin_memoryTrue但不使用non_blockingTrue你只承担了页锁定内存的开销却没有获得异步传输的好处。# 正确pin_memory non_blocking dataloader DataLoader(..., pin_memoryTrue) for x, y in dataloader: x x.cuda(non_blockingTrue) # 异步传输 y y.cuda(non_blockingTrue) # 可以在这里做其他 CPU 操作GPU 在后台传输数据 # 错误只 pin_memory 不 non_blocking for x, y in dataloader: x x.cuda() # 同步传输pin_memory 的优势被浪费3. 内存消耗翻倍pin_memory会在内存中保留一份额外的页锁定副本。如果你的数据集很大或 batch size 很大内存消耗会显著增加# 假设 batch_size128, 图像大小 3x224x224, float32 # 每个 batch 的内存: 128 * 3 * 224 * 224 * 4 ~77MB # pin_memory 会额外占用 77MB 页锁定内存 # 加上 prefetch_factor2可能占用 3x 231MB4. macOS 上的特殊问题在 macOS 上pin_memory的行为可能与 Linux 不同。macOS 的内存管理策略不同页锁定内存的开销可能更大。在 macOS 上建议先进行基准测试再决定是否使用。5. CPU 端预处理瓶颈如果你的数据加载包含大量 CPU 预处理如复杂的图像增强pin_memory的额外 CPU 开销会加剧瓶颈。在这种情况下优化预处理代码比使用pin_memory更重要。6. prefetch_factor 的影响prefetch_factor控制 worker 预加载的 batch 数量。默认为 2。增大prefetch_factor可以更好地重叠数据加载和计算但也会增加内存消耗。7. persistent_workers 的配合persistent_workersTrue可以避免每个 epoch 重新创建 worker 进程的开销。建议与pin_memory一起使用。8. 监控系统资源使用pin_memory时建议监控系统内存使用情况import psutil import os def monitor_memory(): process psutil.Process(os.getpid()) mem_info process.memory_info() print(fRSS: {mem_info.rss / 1024**2:.1f} MB) print(fVMS: {mem_info.vms / 1024**2:.1f} MB) # 系统级内存 vm psutil.virtual_memory() print(fSystem memory: {vm.used / 1024**2:.0f}/{vm.total / 1024**2:.0f} MB ({vm.percent}%))9. 不要在 CPU-only 训练中使用 pin_memorypin_memory的唯一目的是加速 CPU 到 GPU 的传输。如果你在 CPU 上训练pin_memory完全没有用处只会增加开销。10. 使用 PyTorch Profiler 诊断当不确定瓶颈在哪里时使用 PyTorch Profiler 进行精确诊断from torch.profiler import profile, record_function, ProfilerActivity with profile(activities[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapesTrue) as prof: with record_function(model_training): for batch_x, batch_y in dataloader: batch_x batch_x.cuda(non_blockingTrue) batch_y batch_y.cuda(non_blockingTrue) output model(batch_x) loss criterion(output, batch_y) loss.backward() optimizer.step() print(prof.key_averages().table(sort_bycuda_time_total))总结pin_memory在 PyTorch 中是一个有用的但容易被误解的优化技术。它并非银弹在以下条件下才能发挥正面作用大 batch size 大数据当每个 batch 的数据量较大时DMA 传输节省的时间才能超过页锁定内存分配和复制的开销。配合non_blockingTrue这是pin_memory发挥作用的关键。异步传输允许数据传输和 GPU 计算重叠这才是性能提升的核心来源。充足的系统内存页锁定内存是稀缺资源系统内存不足时会导致性能退化。GPU 计算密集型任务当 GPU 计算时间较长时异步传输的重叠效果更明显。在以下情况下pin_memory可能反而变慢小 batch size 或小数据CPU 瓶颈数据预处理耗时系统内存不足未配合non_blockingTrue使用macOS 或某些特殊系统配置最佳实践是始终进行基准测试使用本文提供的诊断工具比较不同配置的性能然后选择最适合你具体场景的配置。不要盲目相信pin_memory 总是更快的说法要根据实际情况做出决策。
【Bug已解决】Pinning memory is actually slower in PyTorch? 解决方案问题描述在 PyTorch 中pin_memory页锁定内存是一个被广泛推荐的数据加载优化技术。官方文档和大量教程都建议在DataLoader中设置pin_memoryTrue以加速从 CPU 到 GPU 的数据传输。然而许多开发者发现在实际使用中开启pin_memory后训练速度不仅没有提升反而变慢了。这种现象看似矛盾——一个本应加速训练的优化手段却导致了性能退化。问题的根源涉及多个层面操作系统内存管理机制、PyTorch 的内存分配策略、数据加载器的工作方式以及硬件配置等。本文将深入分析为什么pin_memory在某些情况下会变慢并提供完整的诊断和解决方案。错误复现错误场景一pin_memory 导致训练变慢import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader import time import numpy as np # 创建一个简单的数据集 class SimpleDataset(Dataset): def __init__(self, num_samples10000, data_size224*224*3): self.data np.random.randn(num_samples, data_size).astype(np.float32) self.labels np.random.randint(0, 10, num_samples) def __len__(self): return len(self.labels) def __getitem__(self, idx): return torch.from_numpy(self.data[idx]), self.labels[idx] dataset SimpleDataset(num_samples10000) # 测试 pin_memoryFalse dataloader_no_pin DataLoader(dataset, batch_size64, shuffleTrue, num_workers4, pin_memoryFalse) # 测试 pin_memoryTrue dataloader_pin DataLoader(dataset, batch_size64, shuffleTrue, num_workers4, pin_memoryTrue) # 简单模型 model nn.Sequential( nn.Linear(224*224*3, 512), nn.ReLU(), nn.Linear(512, 10) ).cuda() criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.parameters()) def train_epoch(dataloader, model, criterion, optimizer): model.train() total_time 0 for batch_x, batch_y in dataloader: start time.time() batch_x batch_x.cuda(non_blockingTrue) batch_y batch_y.cuda(non_blockingTrue) optimizer.zero_grad() output model(batch_x) loss criterion(output, batch_y) loss.backward() optimizer.step() total_time time.time() - start return total_time # 对比测试 print(Testing without pin_memory...) time_no_pin train_epoch(dataloader_no_pin, model, criterion, optimizer) print(fTime without pin_memory: {time_no_pin:.2f}s) print(\nTesting with pin_memory...) time_pin train_epoch(dataloader_pin, model, criterion, optimizer) print(fTime with pin_memory: {time_pin:.2f}s) print(f\nDifference: pin_memory is {slower if time_pin time_no_pin else faster} by {abs(time_pin - time_no_pin):.2f}s)运行结果可能如下Testing without pin_memory... Time without pin_memory: 12.34s Testing with pin_memory... Time with pin_memory: 15.67s Difference: pin_memory is slower by 3.33s错误场景二内存不足导致频繁交换# 当数据集很大时pin_memory 会占用大量页锁定内存 # 导致操作系统频繁进行内存交换swap严重拖慢速度 large_dataset SimpleDataset(num_samples100000, data_size512*512*3) # pin_memoryTrue 会尝试将所有数据复制到页锁定内存 # 如果系统物理内存不足会导致严重的性能问题 dataloader DataLoader(large_dataset, batch_size32, shuffleTrue, num_workers8, pin_memoryTrue) # 系统可能报出内存警告或 OOM try: for batch_x, batch_y in dataloader: batch_x batch_x.cuda() # ... 训练代码 except RuntimeError as e: print(fError: {e}) # 可能的输出: CUDA out of memory 或 系统级别的 OOM错误场景三num_workers 与 pin_memory 的冲突# 在某些系统上多进程 pin_memory 会导致竞争和锁争用 dataloader DataLoader( dataset, batch_size64, shuffleTrue, num_workers16, # 过多的 worker pin_memoryTrue, persistent_workersTrue ) # 在 macOS 或某些 Linux 配置上这可能导致 # 1. worker 进程之间的内存竞争 # 2. 页锁定内存分配失败 # 3. 数据加载变慢而非变快根因分析1. 什么是页锁定内存Pinned Memory在操作系统中内存分为两种可分页内存Pageable Memory普通内存操作系统可以将其换出到磁盘swap space以释放物理内存。页锁定内存Pinned/Page-locked Memory被锁定在物理内存中不会被操作系统换出。GPU 通过 DMADirect Memory Access技术直接从 CPU 内存读取数据。DMA 操作要求源内存是页锁定的因为 DMA 操作不能触发缺页中断。当你将可分页内存传输到 GPU 时驱动程序需要先分配一块页锁定内存将数据从可分页内存复制到页锁定内存通过 DMA 将数据从页锁定内存传输到 GPU释放页锁定内存而使用页锁定内存时步骤 1、2、4 被省略理论上应该更快。2. 为什么 pin_memory 反而变慢原因一内存分配开销pin_memoryTrue时DataLoader 会在主进程中创建一个pin_memory_thread该线程负责将 worker 进程返回的可分页内存数据复制到页锁定内存。这个复制操作本身是有开销的# DataLoader 内部大致做了以下操作 # 1. worker 进程返回可分页内存的 batch # 2. pin_memory_thread 分配页锁定内存 # 3. pin_memory_thread 将数据复制到页锁定内存 # 4. 返回页锁定内存的 batch如果数据量较小或模型计算量较小这个额外的复制开销可能超过 DMA 传输节省的时间。原因二页锁定内存是稀缺资源操作系统的页锁定内存总量是有限的。当pin_memory占用过多页锁定内存时其他系统进程无法使用这些内存操作系统可能被迫进行更多的内存管理操作在内存紧张时系统性能急剧下降原因三与non_blocking的误用pin_memory的真正优势在于配合non_blockingTrue使用时可以实现数据传输和计算的重叠# 正确用法pin_memory non_blocking 实现重叠 batch_x batch_x.cuda(non_blockingTrue) # 异步传输 # 此时 CPU 可以继续执行其他操作GPU 在后台传输数据但如果你的代码中没有使用non_blockingTrue或者数据传输完成后立即需要使用数据没有可重叠的计算pin_memory的优势就无法体现。原因四CPU 瓶颈当训练的瓶颈在 CPU 端数据预处理、数据增强等而非 GPU 端时pin_memory的额外 CPU 开销会加剧瓶颈导致整体变慢。原因五小 batch size 或小数据对于小 batch size 或小尺寸数据DMA 传输本身非常快微秒级页锁定内存的分配和复制开销可能超过传输节省的时间。原因六多 worker 竞争当num_workers较多时多个 worker 进程同时返回数据pin_memory_thread需要串行处理这些数据可能成为瓶颈。同时多个进程竞争页锁定内存分配也可能导致性能下降。3. 性能影响的量化分析因素pin_memory 有利pin_memory 不利Batch Size大 batch128小 batch32数据大小大数据图像、视频小数据文本、表格GPU 利用率GPU 计算密集GPU 空闲等待系统内存充足32GB不足16GBnum_workers少1-4多8non_blocking使用不使用解决方案方案一正确使用 pin_memory non_blockingimport torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader import time class ImageDataset(Dataset): def __init__(self, num_samples5000, size3*224*224): self.data torch.randn(num_samples, size) self.labels torch.randint(0, 10, (num_samples,)) def __len__(self): return len(self.labels) def __getitem__(self, idx): return self.data[idx], self.labels[idx] class SimpleModel(nn.Module): def __init__(self): super().__init__() self.fc nn.Sequential( nn.Linear(3*224*224, 1024), nn.ReLU(), nn.Linear(1024, 512), nn.ReLU(), nn.Linear(512, 10) ) def forward(self, x): return self.fc(x) def benchmark(dataloader, model, use_non_blocking, num_epochs3): model.train() criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.parameters()) total_time 0 for epoch in range(num_epochs): epoch_start time.time() for batch_x, batch_y in dataloader: # 关键使用 non_blockingTrue 实现异步传输 batch_x batch_x.cuda(non_blockinguse_non_blocking) batch_y batch_y.cuda(non_blockinguse_non_blocking) optimizer.zero_grad() output model(batch_x) loss criterion(output, batch_y) loss.backward() optimizer.step() total_time time.time() - epoch_start return total_time / num_epochs # 实验设置 dataset ImageDataset(num_samples5000) model SimpleModel().cuda() configs [ (No pin, no non_blocking, False, False, False), (pin_memory, no non_blocking, True, False, True), (pin_memory non_blocking, True, True, True), ] print(Benchmarking different configurations:\n) for name, pin_mem, non_block, _ in configs: dataloader DataLoader( dataset, batch_size128, shuffleTrue, num_workers4, pin_memorypin_mem ) avg_time benchmark(dataloader, model, non_block) print(f{name}: {avg_time:.2f}s per epoch)方案二根据场景选择是否启用 pin_memoryimport torch from torch.utils.data import DataLoader def create_optimal_dataloader(dataset, batch_size, model_typeimage, data_size_bytesNone, system_memory_gb16): 根据场景自动选择最优的 DataLoader 配置。 Args: dataset: 数据集 batch_size: 批大小 model_type: image图像, text文本, tabular表格 data_size_bytes: 单个样本的字节数 system_memory_gb: 系统物理内存GB Returns: 配置好的 DataLoader # 估算每 batch 的内存占用 if data_size_bytes: batch_memory batch_size * data_size_bytes else: if model_type image: batch_memory batch_size * 3 * 224 * 224 * 4 # float32 elif model_type text: batch_memory batch_size * 512 * 4 else: batch_memory batch_size * 100 * 4 # 决定是否使用 pin_memory # 条件batch 较大、数据较大、系统内存充足 use_pin_memory ( batch_size 64 and model_type image and system_memory_gb 16 and batch_memory system_memory_gb * 1024**3 * 0.3 # 不超过系统内存的30% ) # 决定 num_workers cpu_count torch.get_num_threads() if model_type image: num_workers min(4, cpu_count) else: num_workers min(2, cpu_count) print(fConfiguration:) print(f pin_memory: {use_pin_memory}) print(f num_workers: {num_workers}) print(f Estimated batch memory: {batch_memory / 1024 / 1024:.2f} MB) return DataLoader( dataset, batch_sizebatch_size, shuffleTrue, num_workersnum_workers, pin_memoryuse_pin_memory, persistent_workersnum_workers 0, prefetch_factor2 if num_workers 0 else None ) # 使用示例 # dataloader create_optimal_dataloader(dataset, batch_size128, model_typeimage)方案三手动管理页锁定内存对于需要精细控制的高级用户import torch def manual_pin_transfer(tensor, devicecuda, non_blockingTrue): 手动管理页锁定内存传输。 适用于需要精细控制内存使用的场景。 if not tensor.is_pinned(): # 分配页锁定内存并复制 pinned_tensor tensor.pin_memory() else: pinned_tensor tensor # 异步传输到 GPU return pinned_tensor.to(device, non_blockingnon_blocking) # 使用示例 data torch.randn(64, 3, 224, 224) # 方式1直接传输不使用 pin_memory # data_gpu data.cuda() # 方式2手动 pin 后传输 data_gpu manual_pin_transfer(data, devicecuda, non_blockingTrue) # 检查是否为页锁定内存 print(fIs pinned: {data.is_pinned()}) # False pinned data.pin_memory() print(fIs pinned: {pinned.is_pinned()}) # True完整修复代码以下是一个完整的、可运行的训练脚本包含 pin_memory 性能诊断和优化import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import Dataset, DataLoader import time import numpy as np import os # # 1. 数据集定义 # class SyntheticImageDataset(Dataset): 合成图像数据集用于性能测试。 模拟真实图像分类场景的数据加载。 def __init__(self, num_samples10000, channels3, height224, width224, num_classes10, augmentFalse): self.num_samples num_samples self.channels channels self.height height self.width width self.num_classes num_classes self.augment augment # 不预加载所有数据模拟真实场景 self.labels np.random.randint(0, num_classes, num_samples) def __len__(self): return self.num_samples def __getitem__(self, idx): # 模拟数据加载和预处理 data np.random.randn(self.channels, self.height, self.width).astype(np.float32) if self.augment: # 模拟数据增强开销 data self._augment(data) label self.labels[idx] return torch.from_numpy(data), torch.tensor(label, dtypetorch.long) def _augment(self, data): 简单的数据增强 # 随机水平翻转 if np.random.random() 0.5: data data[:, :, ::-1].copy() # 随机裁剪 h_off np.random.randint(0, 10) w_off np.random.randint(0, 10) data data[:, h_off:h_offself.height-10, w_off:w_offself.width-10] # 重新调整大小简化版 padded np.zeros((self.channels, self.height, self.width), dtypenp.float32) padded[:, :data.shape[1], :data.shape[2]] data return padded # # 2. 模型定义 # class ResNetLikeModel(nn.Module): 简化的 ResNet 风格模型用于性能测试 def __init__(self, num_classes10): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 64, 3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(64, 128, 3, padding1), nn.BatchNorm2d(128), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(128, 256, 3, padding1), nn.BatchNorm2d(256), nn.ReLU(inplaceTrue), nn.AdaptiveAvgPool2d((1, 1)) ) self.classifier nn.Sequential( nn.Flatten(), nn.Linear(256, 128), nn.ReLU(), nn.Dropout(0.5), nn.Linear(128, num_classes) ) def forward(self, x): x self.features(x) x self.classifier(x) return x # # 3. 性能诊断工具 # class DataLoaderBenchmark: DataLoader 性能诊断工具。 帮助确定 pin_memory 是否有益。 def __init__(self, dataset, batch_size64, modelNone, devicecuda): self.dataset dataset self.batch_size batch_size self.device device self.model model or ResNetLikeModel().to(device) self.criterion nn.CrossEntropyLoss() self.optimizer optim.Adam(self.model.parameters()) def _train_epoch(self, dataloader, use_non_blockingTrue): self.model.train() data_times [] compute_times [] total_start time.time() for batch_x, batch_y in dataloader: data_start time.time() batch_x batch_x.to(self.device, non_blockinguse_non_blocking) batch_y batch_y.to(self.device, non_blockinguse_non_blocking) # 等待数据传输完成 if use_non_blocking: torch.cuda.synchronize() data_time time.time() - data_start data_times.append(data_time) compute_start time.time() self.optimizer.zero_grad() output self.model(batch_x) loss self.criterion(output, batch_y) loss.backward() self.optimizer.step() torch.cuda.synchronize() compute_time time.time() - compute_start compute_times.append(compute_time) total_time time.time() - total_start return { total: total_time, data_transfer_avg: np.mean(data_times), compute_avg: np.mean(compute_times), data_transfer_total: np.sum(data_times), compute_total: np.sum(compute_times), } def run_comparison(self, num_workers4, num_epochs3): 对比不同配置的性能。 configs [ {name: No pin_memory, pin_memory: False, non_blocking: False}, {name: pin_memory only, pin_memory: True, non_blocking: False}, {name: pin_memory non_blocking, pin_memory: True, non_blocking: True}, ] results {} for config in configs: print(f\n--- Testing: {config[name]} ---) dataloader DataLoader( self.dataset, batch_sizeself.batch_size, shuffleTrue, num_workersnum_workers, pin_memoryconfig[pin_memory], persistent_workersTrue ) # 预热 print(Warming up...) self._train_epoch(dataloader, config[non_blocking]) # 正式测试 print(Benchmarking...) epoch_results [] for epoch in range(num_epochs): result self._train_epoch(dataloader, config[non_blocking]) epoch_results.append(result) print(f Epoch {epoch1}: {result[total]:.2f}s f(data: {result[data_transfer_avg]*1000:.1f}ms/batch, fcompute: {result[compute_avg]*1000:.1f}ms/batch)) avg_total np.mean([r[total] for r in epoch_results]) avg_data np.mean([r[data_transfer_avg] for r in epoch_results]) avg_compute np.mean([r[compute_avg] for r in epoch_results]) results[config[name]] { avg_total: avg_total, avg_data_transfer: avg_data, avg_compute: avg_compute, } # 打印汇总 print(\n * 60) print(BENCHMARK SUMMARY) print( * 60) for name, result in results.items(): print(f{name:35s}: {result[avg_total]:.2f}s/epoch f(data: {result[avg_data_transfer]*1000:.1f}ms, fcompute: {result[avg_compute]*1000:.1f}ms)) # 推荐配置 fastest min(results.items(), keylambda x: x[1][avg_total]) print(f\nRecommended: {fastest[0]}) return results # # 4. 优化的训练函数 # def optimized_train(model, train_dataset, val_dataset, num_epochs10, devicecuda, batch_size128): 优化的训练函数自动选择最佳 pin_memory 配置。 # 先运行基准测试确定最佳配置 print(Running DataLoader benchmark to determine optimal configuration...) benchmark DataLoaderBenchmark( train_dataset, batch_sizebatch_size, modelmodel, devicedevice ) results benchmark.run_comparison(num_workers4, num_epochs2) # 选择最快的配置 fastest_config min(results.items(), keylambda x: x[1][avg_total]) use_pin_memory pin_memory in fastest_config[0].lower() use_non_blocking non_blocking in fastest_config[0].lower() print(f\nUsing configuration: pin_memory{use_pin_memory}, fnon_blocking{use_non_blocking}) # 创建最优的 DataLoader train_loader DataLoader( train_dataset, batch_sizebatch_size, shuffleTrue, num_workers4, pin_memoryuse_pin_memory, persistent_workersTrue, prefetch_factor2 ) val_loader DataLoader( val_dataset, batch_sizebatch_size, shuffleFalse, num_workers2, pin_memoryuse_pin_memory, persistent_workersTrue ) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.001) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_maxnum_epochs) best_val_acc 0.0 for epoch in range(num_epochs): # 训练 model.train() train_loss 0.0 train_correct 0 train_total 0 epoch_start time.time() for batch_x, batch_y in train_loader: batch_x batch_x.to(device, non_blockinguse_non_blocking) batch_y batch_y.to(device, non_blockinguse_non_blocking) optimizer.zero_grad() output model(batch_x) loss criterion(output, batch_y) loss.backward() optimizer.step() train_loss loss.item() * batch_x.size(0) _, predicted output.max(1) train_correct predicted.eq(batch_y).sum().item() train_total batch_y.size(0) train_loss / train_total train_acc train_correct / train_total epoch_time time.time() - epoch_start # 验证 model.eval() val_loss 0.0 val_correct 0 val_total 0 with torch.no_grad(): for batch_x, batch_y in val_loader: batch_x batch_x.to(device, non_blockinguse_non_blocking) batch_y batch_y.to(device, non_blockinguse_non_blocking) output model(batch_x) loss criterion(output, batch_y) val_loss loss.item() * batch_x.size(0) _, predicted output.max(1) val_correct predicted.eq(batch_y).sum().item() val_total batch_y.size(0) val_loss / val_total val_acc val_correct / val_total scheduler.step() if val_acc best_val_acc: best_val_acc val_acc torch.save(model.state_dict(), best_model.pth) print(fEpoch [{epoch1}/{num_epochs}] ({epoch_time:.1f}s) fTrain Loss: {train_loss:.4f}, Acc: {train_acc:.4f} | fVal Loss: {val_loss:.4f}, Acc: {val_acc:.4f}) return model # # 5. 主程序 # def main(): torch.manual_seed(42) np.random.seed(42) device torch.device(cuda if torch.cuda.is_available() else cpu) print(fDevice: {device}) print(fCPU threads: {torch.get_num_threads()}) if torch.cuda.is_available(): print(fGPU: {torch.cuda.get_device_name()}) print(fGPU Memory: {torch.cuda.get_device_properties(0).total_mem / 1024**3:.1f} GB) # 创建数据集 train_dataset SyntheticImageDataset(num_samples2000, augmentTrue) val_dataset SyntheticImageDataset(num_samples400, augmentFalse) # 创建模型 model ResNetLikeModel(num_classes10).to(device) # 运行优化训练 model optimized_train( model, train_dataset, val_dataset, num_epochs5, devicedevice, batch_size64 ) print(\nTraining complete!) if __name__ __main__: main()常见陷阱与注意事项1. pin_memory 不等于自动加速pin_memoryTrue只是一个可能加速的设置不是保证加速。它的效果取决于数据大小、batch size、系统内存、GPU 利用率等多个因素。盲目开启可能导致性能下降。2. 必须配合 non_blocking 使用pin_memory的真正优势在于配合non_blockingTrue实现数据传输和计算的重叠。如果只设置pin_memoryTrue但不使用non_blockingTrue你只承担了页锁定内存的开销却没有获得异步传输的好处。# 正确pin_memory non_blocking dataloader DataLoader(..., pin_memoryTrue) for x, y in dataloader: x x.cuda(non_blockingTrue) # 异步传输 y y.cuda(non_blockingTrue) # 可以在这里做其他 CPU 操作GPU 在后台传输数据 # 错误只 pin_memory 不 non_blocking for x, y in dataloader: x x.cuda() # 同步传输pin_memory 的优势被浪费3. 内存消耗翻倍pin_memory会在内存中保留一份额外的页锁定副本。如果你的数据集很大或 batch size 很大内存消耗会显著增加# 假设 batch_size128, 图像大小 3x224x224, float32 # 每个 batch 的内存: 128 * 3 * 224 * 224 * 4 ~77MB # pin_memory 会额外占用 77MB 页锁定内存 # 加上 prefetch_factor2可能占用 3x 231MB4. macOS 上的特殊问题在 macOS 上pin_memory的行为可能与 Linux 不同。macOS 的内存管理策略不同页锁定内存的开销可能更大。在 macOS 上建议先进行基准测试再决定是否使用。5. CPU 端预处理瓶颈如果你的数据加载包含大量 CPU 预处理如复杂的图像增强pin_memory的额外 CPU 开销会加剧瓶颈。在这种情况下优化预处理代码比使用pin_memory更重要。6. prefetch_factor 的影响prefetch_factor控制 worker 预加载的 batch 数量。默认为 2。增大prefetch_factor可以更好地重叠数据加载和计算但也会增加内存消耗。7. persistent_workers 的配合persistent_workersTrue可以避免每个 epoch 重新创建 worker 进程的开销。建议与pin_memory一起使用。8. 监控系统资源使用pin_memory时建议监控系统内存使用情况import psutil import os def monitor_memory(): process psutil.Process(os.getpid()) mem_info process.memory_info() print(fRSS: {mem_info.rss / 1024**2:.1f} MB) print(fVMS: {mem_info.vms / 1024**2:.1f} MB) # 系统级内存 vm psutil.virtual_memory() print(fSystem memory: {vm.used / 1024**2:.0f}/{vm.total / 1024**2:.0f} MB ({vm.percent}%))9. 不要在 CPU-only 训练中使用 pin_memorypin_memory的唯一目的是加速 CPU 到 GPU 的传输。如果你在 CPU 上训练pin_memory完全没有用处只会增加开销。10. 使用 PyTorch Profiler 诊断当不确定瓶颈在哪里时使用 PyTorch Profiler 进行精确诊断from torch.profiler import profile, record_function, ProfilerActivity with profile(activities[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapesTrue) as prof: with record_function(model_training): for batch_x, batch_y in dataloader: batch_x batch_x.cuda(non_blockingTrue) batch_y batch_y.cuda(non_blockingTrue) output model(batch_x) loss criterion(output, batch_y) loss.backward() optimizer.step() print(prof.key_averages().table(sort_bycuda_time_total))总结pin_memory在 PyTorch 中是一个有用的但容易被误解的优化技术。它并非银弹在以下条件下才能发挥正面作用大 batch size 大数据当每个 batch 的数据量较大时DMA 传输节省的时间才能超过页锁定内存分配和复制的开销。配合non_blockingTrue这是pin_memory发挥作用的关键。异步传输允许数据传输和 GPU 计算重叠这才是性能提升的核心来源。充足的系统内存页锁定内存是稀缺资源系统内存不足时会导致性能退化。GPU 计算密集型任务当 GPU 计算时间较长时异步传输的重叠效果更明显。在以下情况下pin_memory可能反而变慢小 batch size 或小数据CPU 瓶颈数据预处理耗时系统内存不足未配合non_blockingTrue使用macOS 或某些特殊系统配置最佳实践是始终进行基准测试使用本文提供的诊断工具比较不同配置的性能然后选择最适合你具体场景的配置。不要盲目相信pin_memory 总是更快的说法要根据实际情况做出决策。
