文章目录
- 问题
- 解决方案
问题
在批量图片转向量时,遇到了bad image,说图片不全,需要忽略掉。
解决方案
使用collate_fn
自定义一个方法,核对过滤掉坏数据。个人推测collate_fn
的工作流程如下:
对于每个batch来说,dataloader先用随机采样出一批样本index_lis,然后用dataset.__getitem__函数获取这批样本,再用`collate_fn`函数对获取的样本进行核对整理,最后通过stack操作等将这批数据转为torch.Tensor或者list,numpy等格式的batch数据集,传递给下游。
其具体流程是:
- 在Dataset的
__getitem__
中将坏样本结果定义为None。 - 自定义一下核对函数
my_collate
,去掉结果为None的样本。 - 最后再执行一下
default_collate
函数,这一步很重要,不然就会报错。这个函数,将输入的batch数据,转为Tensor
。如果是tuple元素,那么就转为多个Tensor
。
# 数据集
from torch.utils.data import Dataset, DataLoader, default_collatedef my_collate(batch):# 2. 在`__getitem__`中将坏样本结果定义为空,并过滤掉坏样本batch = [e for e in batch if e]# 3. 一定要最后执行一下这个函数return default_collate(batch)class CustomDataset(Dataset):def __init__(self, img_name_lis, img_path_lis, transform):self.img_name_lis = img_name_lisself.image_path = img_path_lisself.transform = transformdef __len__(self):return len(self.img_name_lis)def __getitem__(self, idx):image_path = self.image_path[idx]image_id = self.img_name_lis[idx]try:img = Image.open(image_path).convert("L").convert("RGB")img = self.transform(img)return image_id, imgexcept Exception as e:print("error image id:", image_id)# 1. 报错了,就忽略并给None结果。return Nonedataset = CustomDataset(image_names, image_path_lis, transformer)
dataloader = DataLoader(dataset,batch_size=550, shuffle=False,pin_memory=True, num_workers=18, collate_fn=my_collate)