【一起來學AI大模型】PyTorch DataLoader 實戰指南

DataLoader 是 PyTorch 中處理數據的核心組件，它提供了高效的數據加載、批處理和并行處理功能。下面是一個全面的 DataLoader 實戰指南，包含代碼示例和最佳實踐。

基礎用法：簡單數據加載

import torch
from torch.utils.data import Dataset, DataLoader# 1. 創建自定義數據集
class SimpleDataset(Dataset):def __init__(self, size=1000):self.data = torch.randn(size, 3, 32, 32)  # 模擬圖像數據self.labels = torch.randint(0, 10, (size,))  # 0-9的標簽def __len__(self):return len(self.data)def __getitem__(self, idx):return self.data[idx], self.labels[idx]# 2. 創建DataLoader
dataset = SimpleDataset(1000)
dataloader = DataLoader(dataset,batch_size=64,       # 批大小shuffle=True,        # 是否打亂數據num_workers=4,       # 使用4個進程加載數據pin_memory=True      # 使用固定內存（加速GPU傳輸）
)# 3. 使用DataLoader
for epoch in range(3):print(f"Epoch {epoch+1}")for batch_idx, (data, targets) in enumerate(dataloader):# 數據自動分批：data.shape = [64, 3, 32, 32], targets.shape = [64]if batch_idx % 10 == 0:print(f"  Batch {batch_idx}: {data.shape}, {targets.shape}")print("Epoch completed\n")

高級功能：自定義數據集與轉換

圖像數據集示例

import os
from PIL import Image
from torchvision import transformsclass CustomImageDataset(Dataset):def __init__(self, img_dir, transform=None):self.img_dir = img_dirself.transform = transformself.img_names = [f for f in os.listdir(img_dir) if f.endswith('.jpg')]# 假設文件名格式為 "label_imageid.jpg"，例如 "3_001.jpg"self.labels = [int(f.split('_')[0]) for f in self.img_names]def __len__(self):return len(self.img_names)def __getitem__(self, idx):img_path = os.path.join(self.img_dir, self.img_names[idx])image = Image.open(img_path).convert('RGB')label = self.labels[idx]if self.transform:image = self.transform(image)return image, label# 定義數據轉換
transform = transforms.Compose([transforms.Resize((256, 256)),      # 調整大小transforms.RandomHorizontalFlip(),   # 隨機水平翻轉transforms.RandomRotation(15),       # 隨機旋轉 ±15度transforms.ToTensor(),               # 轉為Tensor [0,1]transforms.Normalize(                # 標準化mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])# 創建數據集和DataLoader
dataset = CustomImageDataset('/path/to/images', transform=transform)
dataloader = DataLoader(dataset,batch_size=32,shuffle=True,num_workers=4,collate_fn=lambda batch: tuple(zip(*batch))  # 自定義批處理函數
)

文本數據集示例

from torchtext.vocab import build_vocab_from_iterator
from torchtext.data.utils import get_tokenizerclass TextDataset(Dataset):def __init__(self, file_path, max_len=100):self.max_len = max_lenself.tokenizer = get_tokenizer('basic_english')# 讀取文本數據和標簽self.texts = []self.labels = []with open(file_path, 'r', encoding='utf-8') as f:for line in f:label, text = line.split('\t')self.labels.append(int(label))self.texts.append(text.strip())# 構建詞匯表self.vocab = build_vocab_from_iterator((self.tokenizer(text) for text in self.texts),specials=['<unk>', '<pad>'])self.vocab.set_default_index(self.vocab['<unk>'])def __len__(self):return len(self.texts)def __getitem__(self, idx):text = self.texts[idx]tokens = self.tokenizer(text)# 將token轉換為索引indices = [self.vocab[token] for token in tokens]# 截斷或填充序列if len(indices) > self.max_len:indices = indices[:self.max_len]else:indices = indices + [self.vocab['<pad>']] * (self.max_len - len(indices))return torch.tensor(indices), self.labels[idx]# 自定義批處理函數（處理變長序列）
def collate_fn(batch):texts, labels = zip(*batch)# 找到批次中最長序列的長度max_len = max(len(t) for t in texts)# 填充所有序列到相同長度padded_texts = []for text in texts:padding = torch.zeros(max_len - len(text), dtype=torch.long)padded_texts.append(torch.cat((text, padding)))return torch.stack(padded_texts), torch.tensor(labels)# 創建DataLoader
text_dataset = TextDataset('/path/to/text_data.txt', max_len=100)
text_dataloader = DataLoader(text_dataset,batch_size=32,shuffle=True,num_workers=2,collate_fn=collate_fn  # 使用自定義批處理函數
)

性能優化技巧

1. 使用并行加載

# 根據CPU核心數設置num_workers
import os
num_workers = min(4, os.cpu_count())  # 使用不超過4個或CPU核心數的workerdataloader = DataLoader(dataset,batch_size=64,shuffle=True,num_workers=num_workers,pin_memory=True,  # 對于GPU訓練非常重要persistent_workers=True  # 保持worker進程活動（PyTorch 1.7+）
)

2. 數據預取

from torch.utils.data import DataLoader, PrefetchGenerator# 使用預取生成器（PyTorch 1.7+）
dataloader = DataLoader(dataset,batch_size=64,shuffle=True,num_workers=4,prefetch_factor=2  # 每個worker預取的批次數
)# 或者使用自定義預取
class PrefetchLoader:def __init__(self, loader, device):self.loader = loaderself.device = deviceself.stream = torch.cuda.Stream() if device.type == 'cuda' else Nonedef __iter__(self):first = Truefor batch in self.loader:if self.stream is not None:with torch.cuda.stream(self.stream):batch = self._preprocess(batch)else:batch = self._preprocess(batch)if not first and self.stream is not None:torch.cuda.current_stream().wait_stream(self.stream)first = Falseyield batchdef _preprocess(self, batch):data, target = batchreturn data.to(self.device, non_blocking=True), target.to(self.device, non_blocking=True)# 使用自定義預取
device = torch.device('cuda')
prefetch_dataloader = PrefetchLoader(dataloader, device)

3. 內存映射文件處理大文件

import numpy as np
import torch
from torch.utils.data import Datasetclass MmapDataset(Dataset):def __init__(self, file_path, shape, dtype=np.float32):self.data = np.memmap(file_path, dtype=dtype, mode='r', shape=shape)def __len__(self):return self.data.shape[0]def __getitem__(self, idx):return torch.from_numpy(np.array(self.data[idx]))

分布式數據加載

import torch.distributed as dist
from torch.utils.data.distributed import DistributedSampler# 初始化分布式環境
dist.init_process_group(backend='nccl')
rank = dist.get_rank()
world_size = dist.get_world_size()# 創建分布式采樣器
sampler = DistributedSampler(dataset,num_replicas=world_size,rank=rank,shuffle=True,seed=42
)# 創建分布式DataLoader
dist_dataloader = DataLoader(dataset,batch_size=64,sampler=sampler,num_workers=4,pin_memory=True,drop_last=True  # 丟棄最后不完整的批次
)# 在每個進程中
for epoch in range(10):# 設置epoch確保所有進程的shuffle一致dist_dataloader.sampler.set_epoch(epoch)for batch in dist_dataloader:# 處理批次數據pass

數據增強策略

圖像增強

from torchvision import transforms
import albumentations as A
from albumentations.pytorch import ToTensorV2# 使用torchvision
torchvision_transform = transforms.Compose([transforms.RandomResizedCrop(224),transforms.RandomHorizontalFlip(),transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])# 使用Albumentations（更豐富的增強）
albumentations_transform = A.Compose([A.RandomResizedCrop(224, 224),A.HorizontalFlip(p=0.5),A.VerticalFlip(p=0.2),A.Rotate(limit=30),A.RGBShift(r_shift_limit=25, g_shift_limit=25, b_shift_limit=25, p=0.9),A.RandomBrightnessContrast(brightness_limit=0.3, contrast_limit=0.3, p=0.5),A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),ToTensorV2()
])# 在數據集類中使用
def __getitem__(self, idx):img_path = self.img_paths[idx]image = cv2.imread(img_path)image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)if self.transform:augmented = self.transform(image=image)image = augmented['image']return image, self.labels[idx]

文本增強

import nlpaug.augmenter.word as naw# 創建文本增強器
augmenter = naw.ContextualWordEmbsAug(model_path='bert-base-uncased', action="substitute",  # 替換、插入等aug_p=0.1  # 增強比例
)# 在數據集中使用
def __getitem__(self, idx):text = self.texts[idx]if self.augment and random.random() < 0.5:  # 50%概率增強text = augmenter.augment(text)# 后續處理...

數據可視化與調試

import matplotlib.pyplot as plt
import numpy as npdef show_batch(dataloader, n=4):"""顯示一批圖像及其標簽"""dataiter = iter(dataloader)images, labels = next(dataiter)fig, axes = plt.subplots(1, n, figsize=(15, 4))for i in range(n):img = images[i].permute(1, 2, 0).numpy()  # CHW -> HWCimg = img * np.array([0.229, 0.224, 0.225]) + np.array([0.485, 0.456, 0.406])  # 反歸一化img = np.clip(img, 0, 1)axes[i].imshow(img)axes[i].set_title(f"Label: {labels[i].item()}")axes[i].axis('off')plt.show()# 使用
show_batch(dataloader, n=8)

常見問題解決方案

1. 內存不足

# 解決方案1：使用更小的批大小
dataloader = DataLoader(dataset, batch_size=16)# 解決方案2：使用內存映射文件
# 如前文的MmapDataset示例# 解決方案3：使用IterableDataset
from torch.utils.data import IterableDatasetclass LargeIterableDataset(IterableDataset):def __init__(self, file_path, chunk_size=1000):self.file_path = file_pathself.chunk_size = chunk_sizedef __iter__(self):with open(self.file_path, 'r') as f:chunk = []for line in f:chunk.append(process_line(line))  # 自定義處理函數if len(chunk) == self.chunk_size:yield from chunkchunk = []if chunk:yield from chunk# 使用
dataset = LargeIterableDataset('large_file.txt')
dataloader = DataLoader(dataset, batch_size=64)

2. Windows多進程問題

# 解決方案：將主代碼放入if __name__ == '__main__'塊中
if __name__ == '__main__':# 在這里創建DataLoaderdataloader = DataLoader(dataset, num_workers=4)# 訓練代碼...

3. 數據加載成為瓶頸

# 解決方案1：增加num_workers
dataloader = DataLoader(dataset, num_workers=os.cpu_count())# 解決方案2：使用預取
# 如前文的PrefetchLoader示例# 解決方案3：使用更快的存儲（如SSD代替HDD）# 解決方案4：使用更高效的數據格式（如HDF5、LMDB）