Ubuntu下PyTorch数据加载技巧 - 问答

在Ubuntu下使用PyTorch加载数据时，可参考以下技巧提升效率：

一、多线程/多进程加载

使用DataLoader的num_workers参数设置并行加载的子进程数（建议设置为CPU核心数的1-2倍，如num_workers=4），避免数据加载成为瓶颈。
启用pin_memory=True加速数据从CPU到GPU的传输。

二、数据预处理优化

对图像数据，使用高效的预处理库（如turbojpeg）替代默认的PIL，或通过torchvision.transforms进行批量预处理。
确保数据在内存中是连续的，可通过.contiguous()操作优化张量存储。

三、硬件与环境配置

将数据存储在SSD中，避免HDD的I/O延迟。
使用Anaconda管理环境，安装匹配的CUDA和cuDNN版本以加速GPU计算。

四、代码示例

import torch
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms

# 自定义数据集
class CustomDataset(Dataset):
    def __init__(self, data_path, transform=None):
        self.data_path = data_path
        self.transform = transform
        self.image_paths = [...]  # 填充图像路径列表

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        image = Image.open(self.image_paths[idx]).convert('RGB')
        if self.transform:
            image = self.transform(image)
        return image, 0  # 假设标签为0

# 数据加载
transform = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
dataset = CustomDataset(data_path='./data', transform=transform)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4, pin_memory=True)

# 训练循环中使用
for batch in dataloader:
    images, labels = batch
    images, labels = images.cuda(), labels.cuda()  # 移动到GPU
    # 模型训练代码...

五、性能监控

使用nvidia-smi监控GPU使用情况，htop监控CPU和内存，及时调整num_workers避免资源竞争。

0 赞

0 踩