pytorch Dataset,DataLoader产生自定义的训练数据案例
1.torch.utils.data.Dataset
datasets这是一个pytorch定义的dataset的源码集合。下面是一个自定义Datasets的基本框架，初始化放在__init__()中，其中__getitem__()和__len__()两个方法是必须重写的。
__getitem__()返回训练数据，如图片和label，而__len__()返回数据长度。
classCustomDataset(data.Dataset):#需要继承data.Dataset def__init__(self): #TODO #1.Initializefilepathorlistoffilenames. pass def__getitem__(self,index): #TODO #1.Readonedatafromfile(e.g.usingnumpy.fromfile,PIL.Image.open). #2.Preprocessthedata(e.g.torchvision.Transform). #3.Returnadatapair(e.g.imageandlabel). #这里需要注意的是，第一步：readonedata，是一个data pass def__len__(self): #Youshouldchange0tothetotalsizeofyourdataset. return0
2.torch.utils.data.DataLoader
DataLoader(object)可用参数：
dataset(Dataset)传入的数据集
batch_size(int,optional)每个batch有多少个样本
shuffle(bool,optional)在每个epoch开始的时候，对数据进行重新排序
sampler(Sampler,optional)自定义从数据集中取样本的策略，如果指定这个参数，那么shuffle必须为False
batch_sampler(Sampler,optional)与sampler类似，但是一次只返回一个batch的indices（索引），需要注意的是，一旦指定了这个参数，那么batch_size,shuffle,sampler,drop_last就不能再制定了（互斥——Mutuallyexclusive）
num_workers(int,optional)这个参数决定了有几个进程来处理dataloading。0意味着所有的数据都会被load进主进程。（默认为0）
collate_fn(callable,optional)将一个list的sample组成一个mini-batch的函数
pin_memory(bool,optional)如果设置为True，那么dataloader将会在返回它们之前，将tensors拷贝到CUDA中的固定内存（CUDApinnedmemory）中.
drop_last(bool,optional)如果设置为True：这个是对最后的未完成的batch来说的，比如你的batch_size设置为64，而一个epoch只有100个样本，那么训练的时候后面的36个就被扔掉了。如果为False（默认），那么会继续正常执行，只是最后的batch_size会小一点。
timeout(numeric,optional)如果是正数，表明等待从worker进程中收集一个batch等待的时间，若超出设定的时间还没有收集到，那就不收集这个内容了。这个numeric应总是大于等于0。默认为0
worker_init_fn(callable,optional)每个worker初始化函数IfnotNone,thiswillbecalledoneachworkersubprocesswiththeworkerid(anintin[0,num_workers-1])asinput,afterseedingandbeforedataloading.(default:None)
3.使用Dataset,DataLoader产生自定义训练数据
假设TXT文件保存了数据的图片和label,格式如下：第一列是图片的名字，第二列是label
0.jpg0 1.jpg1 2.jpg2 3.jpg3 4.jpg4 5.jpg5 6.jpg6 7.jpg7 8.jpg8 9.jpg9
也可以是多标签的数据，如：
0.jpg010 1.jpg111 2.jpg212 3.jpg313 4.jpg414 5.jpg515 6.jpg616 7.jpg717 8.jpg818 9.jpg919
图库十张原始图片放在./dataset/images目录下，然后我们就可以自定义一个Dataset解析这些数据并读取图片，再使用DataLoader类产生batch的训练数据
3.1自定义Dataset
首先先自定义一个TorchDataset类，用于读取图片数据，产生标签：
注意初始化函数：
importtorch fromtorch.autogradimportVariable fromtorchvisionimporttransforms fromtorch.utils.dataimportDataset,DataLoader importnumpyasnp fromutilsimportimage_processing importos classTorchDataset(Dataset): def__init__(self,filename,image_dir,resize_height=256,resize_width=256,repeat=1): ''' :paramfilename:数据文件TXT：格式：imge_name.jpglabel1_idlabe2_id :paramimage_dir:图片路径：image_dir+imge_name.jpg构成图片的完整路径 :paramresize_height为None时，不进行缩放 :paramresize_width为None时，不进行缩放， PS：当参数resize_height或resize_width其中一个为None时，可实现等比例缩放 :paramrepeat:所有样本数据重复次数，默认循环一次，当repeat为None时，表示无限循环3.2DataLoader产生批训练数据
if__name__=='__main__': train_filename="../dataset/train.txt" #test_filename="../dataset/test.txt" image_dir='../dataset/images' epoch_num=2#总样本循环次数 batch_size=7#训练时的一组数据的大小 train_data_nums=10 max_iterate=int((train_data_nums+batch_size-1)/batch_size*epoch_num)#总迭代次数 train_data=TorchDataset(filename=train_filename,image_dir=image_dir,repeat=1) #test_data=TorchDataset(filename=test_filename,image_dir=image_dir,repeat=1) train_loader=DataLoader(dataset=train_data,batch_size=batch_size,shuffle=False) #test_loader=DataLoader(dataset=test_data,batch_size=batch_size,shuffle=False) #[1]使用epoch方法迭代，TorchDataset的参数repeat=1 forepochinrange(epoch_num): forbatch_image,batch_labelintrain_loader: image=batch_image[0,:] image=image.numpy()#image=np.array(image) image=image.transpose(1,2,0)#通道由[c,h,w]->[h,w,c] image_processing.cv_show_image("image",image) print("batch_image.shape:{},batch_label:{}".format(batch_image.shape,batch_label)) #batch_x,batch_y=Variable(batch_x),Variable(batch_y)上面的迭代代码是通过两个for实现，其中参数epoch_num表示总样本循环次数，比如epoch_num=2，那就是所有样本循环迭代2次。
但这会出现一个问题，当样本总数train_data_nums与batch_size不能整取时，最后一个batch会少于规定batch_size的大小，比如这里样本总数train_data_nums=10，batch_size=7，第一次迭代会产生7个样本，第二次迭代会因为样本不足，只能产生3个样本。
我们希望，每次迭代都会产生相同大小的batch数据，因此可以如下迭代：注意本人在构造TorchDataset类时，就已经考虑循环迭代的方法，因此，你现在只需修改repeat为None时，就表示无限循环了，调用方法如下：
''' 下面两种方式，TorchDataset设置repeat=None可以实现无限循环，退出循环由max_iterate设定 ''' train_data=TorchDataset(filename=train_filename,image_dir=image_dir,repeat=None) train_loader=DataLoader(dataset=train_data,batch_size=batch_size,shuffle=False) #[2]第2种迭代方法 forstep,(batch_image,batch_label)inenumerate(train_loader): image=batch_image[0,:] image=image.numpy()#image=np.array(image) image=image.transpose(1,2,0)#通道由[c,h,w]->[h,w,c] image_processing.cv_show_image("image",image) print("step:{},batch_image.shape:{},batch_label:{}".format(step,batch_image.shape,batch_label)) #batch_x,batch_y=Variable(batch_x),Variable(batch_y) ifstep>=max_iterate: break #[3]第3种迭代方法 #forstepinrange(max_iterate): #batch_image,batch_label=train_loader.__iter__().__next__() #image=batch_image[0,:] #image=image.numpy()#image=np.array(image) #image=image.transpose(1,2,0)#通道由[c,h,w]->[h,w,c] #image_processing.cv_show_image("image",image) #print("batch_image.shape:{},batch_label:{}".format(batch_image.shape,batch_label)) ##batch_x,batch_y=Variable(batch_x),Variable(batch_y)3.3附件：image_processing.py
上面代码，用到image_processing，这是本人封装好的图像处理包，包含读取图片，画图等基本方法：
#-*-coding:utf-8-*- """ @Project:IntelligentManufacture @File:image_processing.py @Author:panjq @E-mail:pan_jinquan@163.com @Date:2019-02-1415:34:50 """ importos importglob importcv2 importnumpyasnp importmatplotlib.pyplotasplt defshow_image(title,image): ''' 调用matplotlib显示RGB图片 :paramtitle:图像标题 :paramimage:图像的数据 :return: ''' #plt.figure("show_image") #print(image.dtype) plt.imshow(image) plt.axis('on')#关掉坐标轴为off plt.title(title)#图像题目 plt.show() defcv_show_image(title,image): ''' 调用OpenCV显示RGB图片 :paramtitle:图像标题 :paramimage:输入RGB图像 :return: ''' channels=image.shape[-1] ifchannels==3: image=cv2.cvtColor(image,cv2.COLOR_RGB2BGR)#将BGR转为RGB cv2.imshow(title,image) cv2.waitKey(0) defread_image(filename,resize_height=None,resize_width=None,normalization=False): ''' 读取图片数据,默认返回的是uint8,[0,255] :paramfilename: :paramresize_height: :paramresize_width: :paramnormalization:是否归一化到[0.,1.0] :return:返回的RGB图片数据 ''' bgr_image=cv2.imread(filename) #bgr_image=cv2.imread(filename,cv2.IMREAD_IGNORE_ORIENTATION|cv2.IMREAD_COLOR) ifbgr_imageisNone: print("Warning:不存在:{}",filename) returnNone iflen(bgr_image.shape)==2:#若是灰度图则转为三通道 print("Warning:grayimage",filename) bgr_image=cv2.cvtColor(bgr_image,cv2.COLOR_GRAY2BGR) rgb_image=cv2.cvtColor(bgr_image,cv2.COLOR_BGR2RGB)#将BGR转为RGB #show_image(filename,rgb_image) #rgb_image=Image.open(filename) rgb_image=resize_image(rgb_image,resize_height,resize_width) rgb_image=np.asanyarray(rgb_image) ifnormalization: #不能写成:rgb_image=rgb_image/255 rgb_image=rgb_image/255.0 #show_image("srcresizeimage",image) returnrgb_image deffast_read_image_roi(filename,orig_rect,ImreadModes=cv2.IMREAD_COLOR,normalization=False): ''' 快速读取图片的方法 :paramfilename:图片路径 :paramorig_rect:原始图片的感兴趣区域rect :paramImreadModes:IMREAD_UNCHANGED IMREAD_GRAYSCALE IMREAD_COLOR IMREAD_ANYDEPTH IMREAD_ANYCOLOR IMREAD_LOAD_GDAL IMREAD_REDUCED_GRAYSCALE_2 IMREAD_REDUCED_COLOR_2 IMREAD_REDUCED_GRAYSCALE_4 IMREAD_REDUCED_COLOR_4 IMREAD_REDUCED_GRAYSCALE_8 IMREAD_REDUCED_COLOR_8 IMREAD_IGNORE_ORIENTATION :paramnormalization:是否归一化 :return:返回感兴趣区域ROI ''' #当采用IMREAD_REDUCED模式时，对应rect也需要缩放 scale=1 ifImreadModes==cv2.IMREAD_REDUCED_COLOR_2orImreadModes==cv2.IMREAD_REDUCED_COLOR_2: scale=1/2 elifImreadModes==cv2.IMREAD_REDUCED_GRAYSCALE_4orImreadModes==cv2.IMREAD_REDUCED_COLOR_4: scale=1/4 elifImreadModes==cv2.IMREAD_REDUCED_GRAYSCALE_8orImreadModes==cv2.IMREAD_REDUCED_COLOR_8: scale=1/8 rect=np.array(orig_rect)*scale rect=rect.astype(int).tolist() bgr_image=cv2.imread(filename,flags=ImreadModes) ifbgr_imageisNone: print("Warning:不存在:{}",filename) returnNone iflen(bgr_image.shape)==3:# rgb_image=cv2.cvtColor(bgr_image,cv2.COLOR_BGR2RGB)#将BGR转为RGB else: rgb_image=bgr_image#若是灰度图 rgb_image=np.asanyarray(rgb_image) ifnormalization: #不能写成:rgb_image=rgb_image/255 rgb_image=rgb_image/255.0 roi_image=get_rect_image(rgb_image,rect) #show_image_rect("srcresizeimage",rgb_image,rect) #cv_show_image("reROI",roi_image) returnroi_image defresize_image(image,resize_height,resize_width): ''' :paramimage: :paramresize_height: :paramresize_width: :return: ''' image_shape=np.shape(image) height=image_shape[0] width=image_shape[1] if(resize_heightisNone)and(resize_widthisNone):#错误写法：resize_heightandresize_widthisNone returnimage ifresize_heightisNone: resize_height=int(height*resize_width/width) elifresize_widthisNone: resize_width=int(width*resize_height/height) image=cv2.resize(image,dsize=(resize_width,resize_height)) returnimage defscale_image(image,scale): ''' :paramimage: :paramscale:(scale_w,scale_h) :return: ''' image=cv2.resize(image,dsize=None,fx=scale[0],fy=scale[1]) returnimage defget_rect_image(image,rect): ''' :paramimage: :paramrect:[x,y,w,h] :return: ''' x,y,w,h=rect cut_img=image[y:(y+h),x:(x+w)] returncut_img defscale_rect(orig_rect,orig_shape,dest_shape): ''' 对图像进行缩放时，对应的rectangle也要进行缩放 :paramorig_rect:原始图像的rect=[x,y,w,h] :paramorig_shape:原始图像的维度shape=[h,w] :paramdest_shape:缩放后图像的维度shape=[h,w] :return:经过缩放后的rectangle ''' new_x=int(orig_rect[0]*dest_shape[1]/orig_shape[1]) new_y=int(orig_rect[1]*dest_shape[0]/orig_shape[0]) new_w=int(orig_rect[2]*dest_shape[1]/orig_shape[1]) new_h=int(orig_rect[3]*dest_shape[0]/orig_shape[0]) dest_rect=[new_x,new_y,new_w,new_h] returndest_rect defshow_image_rect(win_name,image,rect): ''' :paramwin_name: :paramimage: :paramrect: :return: ''' x,y,w,h=rect point1=(x,y) point2=(x+w,y+h) cv2.rectangle(image,point1,point2,(0,0,255),thickness=2) cv_show_image(win_name,image) defrgb_to_gray(image): image=cv2.cvtColor(image,cv2.COLOR_RGB2GRAY) returnimage defsave_image(image_path,rgb_image,toUINT8=True): iftoUINT8: rgb_image=np.asanyarray(rgb_image*255,dtype=np.uint8) iflen(rgb_image.shape)==2:#若是灰度图则转为三通道 bgr_image=cv2.cvtColor(rgb_image,cv2.COLOR_GRAY2BGR) else: bgr_image=cv2.cvtColor(rgb_image,cv2.COLOR_RGB2BGR) cv2.imwrite(image_path,bgr_image) defcombime_save_image(orig_image,dest_image,out_dir,name,prefix): ''' 命名标准：out_dir/name_prefix.jpg :paramorig_image: :paramdest_image: :paramimage_path: :paramout_dir: :paramprefix: :return: ''' dest_path=os.path.join(out_dir,name+"_"+prefix+".jpg") save_image(dest_path,dest_image) dest_image=np.hstack((orig_image,dest_image)) save_image(os.path.join(out_dir,"{}_src_{}.jpg".format(name,prefix)),dest_image)3.4完整的代码
#-*-coding:utf-8-*- """ @Project:pytorch-learning-tutorials @File:dataset.py @Author:panjq @E-mail:pan_jinquan@163.com @Date:2019-03-0718:45:06 """ importtorch fromtorch.autogradimportVariable fromtorchvisionimporttransforms fromtorch.utils.dataimportDataset,DataLoader importnumpyasnp fromutilsimportimage_processing importos classTorchDataset(Dataset): def__init__(self,filename,image_dir,resize_height=256,resize_width=256,repeat=1): ''' :paramfilename:数据文件TXT：格式：imge_name.jpglabel1_idlabe2_id :paramimage_dir:图片路径：image_dir+imge_name.jpg构成图片的完整路径 :paramresize_height为None时，不进行缩放 :paramresize_width为None时，不进行缩放， PS：当参数resize_height或resize_width其中一个为None时，可实现等比例缩放 :paramrepeat:所有样本数据重复次数，默认循环一次，当repeat为None时，表示无限循环[h,w,c] image_processing.cv_show_image("image",image) print("batch_image.shape:{},batch_label:{}".format(batch_image.shape,batch_label)) #batch_x,batch_y=Variable(batch_x),Variable(batch_y) ''' 下面两种方式，TorchDataset设置repeat=None可以实现无限循环，退出循环由max_iterate设定 ''' train_data=TorchDataset(filename=train_filename,image_dir=image_dir,repeat=None) train_loader=DataLoader(dataset=train_data,batch_size=batch_size,shuffle=False) #[2]第2种迭代方法 forstep,(batch_image,batch_label)inenumerate(train_loader): image=batch_image[0,:] image=image.numpy()#image=np.array(image) image=image.transpose(1,2,0)#通道由[c,h,w]->[h,w,c] image_processing.cv_show_image("image",image) print("step:{},batch_image.shape:{},batch_label:{}".format(step,batch_image.shape,batch_label)) #batch_x,batch_y=Variable(batch_x),Variable(batch_y) ifstep>=max_iterate: break #[3]第3种迭代方法 #forstepinrange(max_iterate): #batch_image,batch_label=train_loader.__iter__().__next__() #image=batch_image[0,:] #image=image.numpy()#image=np.array(image) #image=image.transpose(1,2,0)#通道由[c,h,w]->[h,w,c] #image_processing.cv_show_image("image",image) #print("batch_image.shape:{},batch_label:{}".format(batch_image.shape,batch_label)) ##batch_x,batch_y=Variable(batch_x),Variable(batch_y) 以上为个人经验，希望能给大家一个参考，也希望大家多多支持毛票票。如有错误或未考虑完全的地方，望不吝赐教。
声明：本文内容来源于网络，版权归原作者所有，内容由互联网用户自发贡献自行上传，本网站不拥有所有权，未作人工编辑处理，也不承担相关法律责任。如果您发现有涉嫌版权的内容，欢迎发送邮件至：czq8825#qq.com（发邮件时，请将#更换为@）进行举报，并提供相关证据，一经查实，本站将立刻删除涉嫌侵权内容。