Pytorch读取数据Dataset,DataLoader及流式读取文件

最新推荐文章于 2025-07-11 10:04:38 发布

rosefunR

最新推荐文章于 2025-07-11 10:04:38 发布

阅读量8.4k

点赞数 5

CC 4.0 BY-SA版权

分类专栏：学习Pytorch

本文为博主CSDN Rosefun96原创文章，转载请注明。

本文链接：https://blog.csdn.net/rosefun96/article/details/87947590

简介

最近都是看图像里边的语义分割部分内容，比较有趣，同时入门Pytorch。Pytorch的主要特点是基本上所有操作都是用类来进行封装，本身自带很多类，而且你也可以根据官方的类进行修改。

数据导入，本来Pytorch就有好几个类共同实现，分别是 DataSet, DataLoader, DataLoaderIter等。
DataSet指明读取图片的路径和数目，DataLoader 实现读取数据的并行。

1. ImageFolder + DataLoader

首先我的数据是存在data_dir里边，每个子文件夹作为一类。

data_dir = '/Ryoma/data/'
from torchvision import transforms

transform = transforms.Compose([
    # you can add other transformations in this list
    transforms.ToTensor()
])

train_sets = datasets.ImageFolder(data_dir, transform)
train_loader = torch.utils.data.DataLoader(train_sets, batch_size=10, 
                                           shuffle=True, num_workers=4)
print(train_loader)
inputs, classes = next(iter(train_loader))

# Visualize a few images
def imshow(inp, title=None):
    """Imshow for Tensor."""
    print(inputs.shape)
    inp = inp[0]
    inp = inp.numpy().transpose((1, 2, 0))
#     mean = np.array([0.485, 0.456, 0.406])
#     std = np.array([0.229, 0.224, 0.225])
#     inp = std * inp + mean
    plt.imshow(inp)
    if title is not None:
        plt.title(title)

imshow(inputs)

划分数据集
如果需要对数据集进行划分，可以采用以下方法：

   num_train = len(train_dataset)
   indices = list(range(num_train))
   split = int(np.floor(valid_size * num_train))

   if shuffle:
       np.random.seed(random_seed)
       np.random.shuffle(indices)

   train_idx, valid_idx = indices[split:], indices[:split]
   train_sampler = SubsetRandomSampler(train_idx)
   valid_sampler = SubsetRandomSampler(valid_idx)

   train_loader = torch.utils.data.DataLoader(
       train_dataset, batch_size=batch_size, sampler=train_sampler,
       num_workers=num_workers, pin_memory=pin_memory,
   )
   vali