通过pin_memory 优化 PyTorch 数据加载和传输：工作原理、使用场景与性能分析

import torch
 from torch.utils.data import DataLoader
 from torchvision import datasets, transforms
 
 # 定义图像转换
 transform = transforms.Compose([
     transforms.Resize((256, 256)),
     transforms.ToTensor(),
 ])
 
 # 加载数据集
 dataset = datasets.ImageFolder(root='path/to/data', transform=transform)
 
 # 使用 pin_memory=True 的 DataLoader  
 dataloader = DataLoader(
     dataset,
     batch_size=64,
     shuffle=True,
     num_workers=4,
     pin_memory=True  # 加快数据向 GPU 的传输速度  
 )
 
 # 数据传输至 GPU
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 for batch in dataloader:
     images, labels = batch
     images = images.to(device, non_blocking=True)  # 更快的传输
     # 训练循环代码

import torch
 from torch import nn  
 from torch.utils.data import DataLoader
 from torchvision import datasets, transforms
 
 # 多 GPU 设置
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 # 示例网络和 DataLoader 设置
 model = nn.DataParallel(nn.Linear(256*256*3, 10)).to(device)
 dataloader = DataLoader(
     datasets.ImageFolder('path/to/data', transform=transform),
     batch_size=64,
     shuffle=True,
     num_workers=4,
     pin_memory=True  # 为跨 GPU 的快速传输启用
 )
 
 # 多 GPU 训练循环
 for batch in dataloader:
     inputs, targets = batch
     inputs, targets = inputs.to(device, non_blocking=True), targets.to(device)
     outputs = model(inputs)
     # 其他训练步骤

import torch  
 from torchvision import transforms
 from PIL import Image
 
 # 定义实时推理的图像转换
 transform = transforms.Compose([
     transforms.Resize((256, 256)),
     transforms.ToTensor()
 ])
 
 # 加载图像并固定内存  
 def load_image(image_path):
     image = Image.open(image_path)
     image = transform(image).unsqueeze(0)
     return image.pin_memory()  # 为推理显式固定内存
 
 # 实时推理
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  
 model = torch.load('model.pth').to(device).eval()
 
 def infer(image_path):
     image = load_image(image_path)
     with torch.no_grad():
         image = image.to(device, non_blocking=True)
         output = model(image)
     return output
 
 # 运行推理
 output = infer('path/to/image.jpg')
 print("推理结果:", output)

import torch
 from torch.utils.data import DataLoader, TensorDataset
 
 # 小数据集示例
 data = torch.randn(100, 10)  # 100 个样本, 10 个特征
 labels = torch.randint(0, 2, (100,))
 
 # 数据集和 DataLoader 设置
 dataset = TensorDataset(data, labels)  
 dataloader = DataLoader(dataset, batch_size=10, shuffle=True, pin_memory=True)
 
 # 简单的基于 CPU 的模型
 model = torch.nn.Linear(10, 2)
 
 # 在 CPU 上的训练循环  
 device = torch.device('cpu')
 for batch_data, batch_labels in dataloader:
     batch_data, batch_labels = batch_data.to(device), batch_labels.to(device)
     # 前向传递
     outputs = model(batch_data)
     # 执行其他训练步骤

import torch
 from torch.utils.data import DataLoader, TensorDataset  
 import torch.utils.benchmark as benchmark
 
 # 用于基准测试的大型数据集
 data = torch.randn(10000, 256)
 labels = torch.randint(0, 10, (10000,))
 dataset = TensorDataset(data, labels)
 
 # 使用 pin_memory=True 和 pin_memory=False 进行基准测试
 def benchmark_loader(pin_memory):
     dataloader = DataLoader(dataset, batch_size=128, pin_memory=pin_memory)
     device = torch.device('cuda')
         
     def load_batch():
         for batch_data, _ in dataloader:
             batch_data = batch_data.to(device, non_blocking=True)
 
     return benchmark.Timer(stmt="load_batch()", globals={"load_batch": load_batch}).timeit(10)
 
 # 结果  
 time_with_pin_memory = benchmark_loader(pin_memory=True)
 time_without_pin_memory = benchmark_loader(pin_memory=False)
 
 print(f"使用 pin_memory=True 的时间: {time_with_pin_memory}")
 print(f"使用 pin_memory=False 的时间: {time_without_pin_memory}")

import torch
 from torch.utils.data import DataLoader, TensorDataset
 from torch.autograd import profiler
 
 # 示例数据集
 data = torch.randn(10000, 256)  
 labels = torch.randint(0, 10, (10000,))
 dataset = TensorDataset(data, labels)
 dataloader = DataLoader(dataset, batch_size=128, pin_memory=True)
 
 # 使用 pin_memory=True 分析数据传输
 device = torch.device('cuda')
 
 def load_and_transfer():
     for batch_data, _ in dataloader:
         batch_data = batch_data.to(device, non_blocking=True)
 
 with profiler.profile(record_shapes=True) as prof:  
     load_and_transfer()
 
 # 显示分析结果
 print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))