上接《关于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)