一个简单的不需要重新定义dataset的方法

上接《关于VGG19训练8分类模型的总结》

其实在文件储存格式为data\train\...jpg的数据集导入时，可以直接用pytorch自带的一个函数导入。而不需要重写dataset的文件。我之前自己写自定义dataset的文件的时候直接被数据类型的报错搞疯，后来才发现竟然可以直接一步到位。贼气。

再重申一遍，文件层次是

data

        train

                cat(这是你分类的种类)

                        0.jpg

                        1.jpg

                         ......

                dog

        validation

                cat

                        .......

                dog    

data_dir = "./dam-v4"##我这里的文件地址格式是./dam-v4/train(validation)/class_name/.jpg
input_shape = 224

batch_size = 32
#data transformation
data_transforms = {
   'train': transforms.Compose([
       transforms.CenterCrop(input_shape),
       transforms.ToTensor(),

   ]),
   'validation': transforms.Compose([
       transforms.CenterCrop(input_shape),
       transforms.ToTensor(),

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

 
class_names = image_datasets['train'].classes
print("数据读取成功")
print(dataset_sizes)