在前面,我们使用Lenet训练的都是使用默认数据加载器加载特定的数据,本章节我们分析下怎么使用自定义的data.Dataset加载数据
口罩数据集
数据分为两类,mask和no_mask,数据集全部来自于网络
- mask数据
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-fNt7ei6M-1610790347001)(https://i.loli.net/2021/01/16/h1TvjaeJ6QzOCR5.png)]
- no_mask数据
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-GFz7JWpW-1610790347003)(https://i.loli.net/2021/01/16/KlawhSNpDURXFEW.png)]
数据目录结构为:
├── mask
│ ├── test
│ │ ├── mask_0
│ │ └── no_mask_1
│ └── train
│ ├── mask_0
│ └── no_mask_1
数据集下载
链接:https://pan.baidu.com/s/10UcSznNbaUJn8EsVeQc2yg
提取码:2kjb
自定义数据集加载器
主要思路是,通过分割train,和test下各个目录下的图片目录,解析出分类名称和id,例如
mask_0 ---> class:mask class_id:0
no_mask_1 ---> class:no_mask class_id:1
关键函数如下
find_classes
def find_classes_with_id(dir:str) -> Tuple[List[str], Dict[str, int]]:
classes = []
class_ids = {}
class_dirs = [d.name for d in os.scandir(dir) if d.is_dir]
for class_dir in class_dirs:
split_list = class_dir.split('_')
if len(split_list)==1:
msg = "{} form is not right, it should be [classname_id]!!\n".format(class_dir)
raise RuntimeError(msg)
class_id = split_list[len(split_list)-1]
if not class_id.isdigit():
msg = "{} is not end with '_digit' !\n".format(class_dir)
raise RuntimeError(msg)
finaly_split_s = '_'+class_id
class_name = class_dir.split(finaly_split_s)[0]
classes.append(class_name)
class_ids[class_name] = class_id
return classes, class_ids
返回的结果如下
['no_mask', 'mask']
{'no_mask': '1', 'mask': '0'}
make_dataset
def make_dataset(
directory: str,
class_to_idx: Dict[str, int],) -> List[Tuple[str, int]]:
instances = []#struct
if not os.path.isdir(directory):
raise ValueError("Image not dir!!!")
image_count = 0
for target_class in sorted(class_to_idx.keys()):
class_index = class_to_idx[target_class]
target_class_dir = target_class+'_'+str(class_index)
target_dir = os.path.join(directory, target_class_dir)
if not os.path.isdir(target_dir):
continue
for root, _, fnames in sorted(os.walk(target_dir, followlinks=True)):
for fname in sorted(fnames):
path = os.path.join(root, fname)
if is_valid_file(path):
item = path, class_index
instances.append(item)
return instances
解析每一个分类目录下图片,返回每个图片的路径和分类id结果,每个item的格式如下
('./datas/mask/train/mask_0/08021120510000090.jpg', 0)
定义数据加载器 CommonData
然后,集成继承data.Dataset实现CommonData加载器
class CommonData(data.Dataset):
def __init__(self, root_dir, transform=None):
self.root_dir = root_dir
self.transform = transform
self.mean = (0.479, 0.385, 0.352)
self.std = (0.194, 0.171, 0.165)
#获取类别和类别id
classes, class_to_idx = find_classes_with_id(root_dir)
samples = make_dataset(root_dir, class_to_idx)
if len(samples) == 0:
msg = "Found 0 files in subfolders of: {}\n".format(root_dir)
raise RuntimeError(msg)
self.classes = classes
self.class_to_idx = class_to_idx
self.samples = samples
self.targets = [s[1] for s in samples]
print('------SmokeData[%s]',root_dir)
print('classes[%s]'% self.classes)
print('class_to_idx[%s]'%self.class_to_idx)
self.count = 0
#print('targets[%s]'%self.targets)
def __len__(self):
return len(self.samples)
def __getitem__(self, idx):
path, target = self.samples[idx]
sample = default_loader(path)
self.count = self.count+1
#sample.save("./"+str(self.count)+'.jpg')
if self.transform is not None:
sample = self.transform(sample)
return sample, target
主要实现__len__和__getitem__方法,分别返回数据集的长度,和遍历获取每个数据的图片和分类id。
使用数据加载器
def getClassfierDataset(train_dir, test_dir, dataresize=64):
train_transforms = transforms.Compose([
transforms.RandomRotation(20),
transforms.Resize((dataresize,dataresize)),
transforms.RandomHorizontalFlip(0.5),
#transforms.ColorJitter(brightness=[0.8,1.3], contrast=[0.8,1.3], saturation=[0.8,1.3], hue=0.2),
transforms.ToTensor(),
transforms.Normalize((0.479, 0.385, 0.352),
(0.194, 0.171, 0.165))])
test_transforms = transforms.Compose([
transforms.Resize((dataresize,dataresize)),
transforms.ToTensor(),
transforms.Normalize((0.479, 0.385, 0.352),
(0.194, 0.171, 0.165))])
tain_smoke_data = CommonData(train_dir, transform = train_transforms)
test_smoke_data = CommonData(test_dir, transform = test_transforms)
# 使用预处理格式加载图像
#train_data = datasets.ImageFolder(train_dir,transform = train_transforms)
#valid_data = datasets.ImageFolder(test_dir,transform = test_transforms)
# 创建三个加载器,分别为训练,验证,测试,将训练集的batch大小设为64,即每次加载器向网络输送64张图片
#shuffle 随机打乱,网络更容易学习不同的特征,更容易收敛
print('load dataset......')
trainloader = torch.utils.data.DataLoader(tain_smoke_data,batch_size = 64,shuffle = True)
validloader = torch.utils.data.DataLoader(test_smoke_data,batch_size = 64)
return trainloader,validloader
通过以上步骤,我们就获取了训练集和验证集的数据加载器,然后训练的时候使用方法如下
train_loader,test_loader=getClassfierDataset(train_dir, test_dir, self.input_shape)
以上就是自定义pytroch数据加载器的具体实现