【PyTorch学习】（三）自定义Datasets

最新推荐文章于 2025-10-10 18:28:14 发布

原创

最新推荐文章于 2025-10-10 18:28:14 发布 · 3.7k 阅读

29 ·

CC 4.0 BY-SA版权

文章标签：

#PyTorch #datasets

本文介绍了PyTorch中数据处理的重要组成部分——Datasets和Dataloader，特别是如何根据实际需求自定义Datasets。以torchvision.datasets.ImageFolder为例，说明其要求的数据结构，并指出在目标检测等场景中，此结构不适用。文章详细讲解了如何重载`__len__`和`__getitem__`函数来自定义Dataset，并展示了CocoDetection dataset的实现方式。最后，提供了一个自定义FaceLandmarksDataset的例子，用于从CSV文件加载人脸关键点坐标。

torchvision.datasets源码地址:https://github.com/pytorch/vision/blob/master/torchvision/datasets

前两篇从搭建经典的ResNet，DenseNet入手简单的了解了下PyTorch搭建网络的方式，但训练一个模型光光搭建好一个网络是不够的，正所谓巧妇难为无米之炊，如何将数据处理成网络可以传递的Tensor也尤为重要，而数据准备过程最最最最最重要的就是Datasets和Dataloader两部分！

torchvision.datasets.ImageFolder就是官方给出的一个datasets的事例，具体使用直接贴上官方tutorial上的代码供参考：

data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

data_dir = 'hymenoptera_data'
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
                                          data_transforms[x])
                  for x in ['train', 'val']}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4,
                                             shuffle=True, num_workers=4)

最低0.47元/天解锁文章