提升PyTorch DataLoader效率：避免重复实例化
在PyTorch深度学习训练中，高效的数据加载至关重要。 反复创建DataLoader实例会导致进程池的重复创建和销毁，严重影响训练速度。本文介绍如何复用DataLoader，避免这种低效的重复实例化操作。
问题：许多代码在每次迭代中都重新创建DataLoader：DataLoader(dataset, batch_size=batch_size, num_workers=num_workers)。 这会造成性能瓶颈，因为DataLoader初始化需要创建进程池，频繁地创建和销毁进程池会消耗大量资源。
AI舌诊面诊问诊医疗系统+诊疗商城系统源码
下载
AI 舌诊面诊问诊医疗系统 + 诊疗商城系统源码是一套结合 AI 辅助诊断与商品转化的健康类程序，基于浩创云、ThinkPHP（PHP）开发，支持微信服务号与 H5 端，适合做中医舌诊面诊问答、健康建议推送以及配套诊疗商品的运营。
解决方案：将DataLoader的创建移至训练循环之外。 只需在训练开始前创建一次DataLoader实例，并在训练循环中重复使用它即可。 以下代码演示了改进后的方法：
import torch
from torch.utils.data import DataLoader, Dataset
from math import sqrt
from typing import List, Tuple, Union
from numpy import ndarray
from PIL import Image
from torchvision import transforms
preprocess = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]
)
])
class PreprocessImageDataset(Dataset):
def __init__(self, images: Union[List[ndarray], Tuple[ndarray]]):
self.images = images
def __len__(self):
return len(self.images)
def __getitem__(self, idx):
image = self.images[idx]
image = Image.fromarray(image)
preprocessed_image: torch.Tensor = preprocess(image)
unsqueezed_image = preprocessed_image
return unsqueezed_image
if __name__=='__main__':
data = list(range(10000000))
batch_size = 10
num_workers = 16
dataset = PreprocessImageDataset(data)
dataloader = DataLoader(dataset, batch_size=batch_size,
num_workers=num_workers)
for epoch in range(5):
print(f"Epoch {epoch + 1}:")
for batch_data in dataloader:
batch_data
print("Batch data:", batch_data)
print("Batch data type :", type(batch_data))
print("Batch data shape:", batch_data.shape)
通过将DataLoader的实例化放在循环外，并在多个epoch中复用同一个实例，我们避免了重复创建进程池，显著提高了数据加载效率，减少了系统开销，从而提升了训练性能。
