PyTorch提供了ImageFolder的类来加载文件结构如下的图片数据集：

root/dog/xxx.pngroot/dog/xxy.pngroot/dog/xxz.pngroot/cat/123.pngroot/cat/nsdf3.pngroot/cat/asd932_.png

使用这个类的问题在于无法将训练集（training dataset）和验证集（validation dataset）分开。我写了两个类来完成这个工作。

import osimport torchfrom torch.utils.data import Dataset, DataLoaderfrom torchvision.transforms import ToTensor, Resize, Composefrom PIL import Imagefrom sklearn.model_selection import train_test_splitclass ImageFolderSplitter: # images should be placed in folders like: # --root # ----root\dogs # ----root\dogs\image1.png # ----root\dogs\image2.png # ----root\cats # ----root\cats\image1.png # ----root\cats\image2.png # path: the root of the image folder def __init__(self, path, train_size = 0.8): self.path = path self.train_size = train_size self.class2num = {} self.num2class = {} self.class_nums = {} self.data_x_path = [] self.data_y_label = [] self.x_train = [] self.x_valid = [] self.y_train = [] self.y_valid = [] for root, dirs, files in os.walk(path): if len(files) == 0 and len(dirs) > 1: for i, dir1 in enumerate(dirs): self.num2class[i] = dir1 self.class2num[dir1] = i elif len(files) > 1 and len(dirs) == 0: category = "" for key in self.class2num.keys(): if key in root: category = key break label = self.class2num[category] self.class_nums[label] = 0 for file1 in files: self.data_x_path.append(os.path.join(root, file1)) self.data_y_label.append(label) self.class_nums[label] += 1 else: raise RuntimeError("please check the folder structure!") self.x_train, self.x_valid, self.y_train, self.y_valid = train_test_split(self.data_x_path, self.data_y_label, shuffle = True, train_size = self.train_size) def getTrainingDataset(self): return self.x_train, self.y_train def getValidationDataset(self): return self.x_valid, self.y_validclass DatasetFromFilename(Dataset): # x: a list of image file full path # y: a list of image categories def __init__(self, x, y, transforms = None): super(DatasetFromFilename, self).__init__() self.x = x self.y = y if transforms == None: self.transforms = ToTensor() else: self.transforms = transforms def __len__(self): return len(self.x) def __getitem__(self, idx): img = Image.open(self.x[idx]) img = img.convert("RGB") return self.transforms(img), torch.tensor([[self.y[idx]]])# test code# splitter = ImageFolderSplitter("for_test")# transforms = Compose([Resize((51, 51)), ToTensor()])# x_train, y_train = splitter.getTrainingDataset()# training_dataset = DatasetFromFilename(x_train, y_train, transforms=transforms)# training_dataloader = DataLoader(training_dataset, batch_size=2, shuffle=True)# x_valid, y_valid = splitter.getValidationDataset()# validation_dataset = DatasetFromFilename(x_valid, y_valid, transforms=transforms)# validation_dataloader = DataLoader(validation_dataset, batch_size=2, shuffle=True)# for x, y in training_dataloader:# print(x.shape, y.shape)

更多的代码可以在我的Github reop下找到。