1. torch.utils.data.Dataset [读取数据集]
datasets这是一个pytorch定义的dataset的源码集合。下面是一个自定义Datasets的基本框架，初始化放在__init__()中，其中__getitem__()和__len__()两个方法是必须重写的。__getitem__()返回训练数据，如图片和label，而__len__()返回数据长度。
class CustomDataset(data.Dataset):#需要继承data.Dataset
def __init__(self):
# TODO
# 1. Initialize file path or list of file names.
# 初始化文件路径或者给出一系列文件名(train、test、val等数据文件名)
pass
def __getitem__(self, index):
# TODO
# 1. Read one data from file (e.g. using numpy.fromfile, PIL.Image.open).
# 2. Preprocess the data (e.g. torchvision.Transform).
# 3. Return a data pair (e.g. image and label).
# 根据索引index取出记录，也可传入文件名以区分train和test，从而返回含或不含label的记录
#这里需要注意的是，第一步：read one data，是一个data
pass
def __len__(self):
# You should change 0 to the total size of your dataset.
# 返回data_len
return 0
复制代码2. torch.utils.data.DataLoader [产生batch数据]
DataLoader(object)可用参数如下：
参数：
dataset(Dataset): 传入的数据集batch_size(int, optional): 每个batch有多少个样本shuffle(bool, optional): 在每个epoch开始的时候，对数据进行重新排序sampler(Sampler, optional): 自定义从数据集中取样本的策略，如果指定这个参数，那么shuffle必须为Falsebatch_sampler(Sampler, optional): 与sampler类似，但是一次只返回一个batch的indices（索引），需要注意的是，一旦指定了这个参数，那么batch_size,shuffle,sampler,drop_last就不能再制定了（互斥——Mutually exclusive）num_workers(int, optional): 这个参数决定了有几个进程来处理data loading。0意味着所有的数据都会被load进主进程。（默认为0）collate_fn(callable, optional): 将一个list的sample组成一个mini-batch的函数pin_memory(bool, optional)： 如果设置为True，那么data loader将会在返回它们之前，将tensors拷贝到CUDA中的固定内存（CUDA pinned memory）中.- **
drop_last(bool, optional):**如果设置为True：这个是对最后的未完成的batch来说的，比如你的batch_size设置为64，而一个epoch只有100个样本，那么训练的时候后面的36个就被扔掉了。 如果为False（默认），那么会继续正常执行，只是最后的batch_size会小一点。- **
timeout(numeric, optional):**如果是正数，表明等待从worker进程中收集一个batch等待的时间，若超出设定的时间还没有收集到，那就不收集这个内容了。这个numeric应总是大于等于0。默认为0worker_init_fn(callable, optional): 每个worker初始化函数 If not None, this will be called on eachworker subprocess with the worker id (an int in [0, num_workers - 1]) as input, after seeding and before data loading. (default: None)
3. 使用Dataset, DataLoader产生自定义训练数据
假设TXT文件保存了数据的图片和label,格式如下：第一列是图片的名字，第二列是label
0.jpg 0
1.jpg 1
2.jpg 2
3.jpg 3
4.jpg 4
5.jpg 5
6.jpg 6
7.jpg 7
8.jpg 8
9.jpg 9
复制代码也可以是多标签的数据，如：
0.jpg 0 10
1.jpg 1 11
2.jpg 2 12
3.jpg 3 13
4.jpg 4 14
5.jpg 5 15
6.jpg 6 16
7.jpg 7 17
8.jpg 8 18
9.jpg 9 19
复制代码图库十张原始图片放在./dataset/images目录下，然后我们就可以自定义一个Dataset解析这些数据并读取图片，再使用DataLoader类产生batch的训练数据
3.1 自定义Dataset
首先先自定义一个TorchDataset类，用于读取图片数据，产生标签：
注意初始化函数：
import torch
from torch.autograd import Variable
from torchvision import transforms
from torch.utils.data import Dataset, DataLoader
import numpy as np
from utils import image_processing
import os
class TorchDataset(Dataset):
def __init__(self, filename, image_dir, resize_height=256, resize_width=256, repeat=1):
'''
:param filename: 数据文件TXT：格式：imge_name.jpg label1_id labe2_id
:param image_dir: 图片路径：image_dir+imge_name.jpg构成图片的完整路径
:param resize_height 为None时，不进行缩放
:param resize_width 为None时，不进行缩放，
PS：当参数resize_height或resize_width其中一个为None时，可实现等比例缩放
:param repeat: 所有样本数据重复次数，默认循环一次，当repeat为None时，表示无限循环<sys.maxsize
'''
self.image_label_list = self.read_file(filename)
self.image_dir = image_dir
self.len = len(self.image_label_list)
self.repeat = repeat
self.resize_height = resize_height
self.resize_width = resize_width
# 相关预处理的初始化
'''class torchvision.transforms.ToTensor'''
# 把shape=(H,W,C)的像素值范围为[0, 255]的PIL.Image或者numpy.ndarray数据
# 转换成shape=(C,H,W)的像素数据，并且被归一化到[0.0, 1.0]的torch.FloatTensor类型。
self.toTensor = transforms.ToTensor()
'''class torchvision.transforms.Normalize(mean, std)
此转换类作用于torch. * Tensor,给定均值(R, G, B) 和标准差(R, G, B)，
用公式channel = (channel - mean) / std进行规范化。
'''
# self.normalize=transforms.Normalize()
def __getitem__(self, i):
index = i % self.len
# print("i={},index={}".format(i, index))
image_name, label = self.image_label_list[index]
image_path = os.path.join(self.image_dir, image_name)
img = self.load_data(image_path, self.resize_height, self.resize_width, normalization=False)
img = self.data_preproccess(img)
label=np.array(label)
return img, label
def __len__(self):
if self.repeat == None:
data_len = 10000000
else:
data_len = len(self.image_label_list) * self.repeat
return data_len
def read_file(self, filename):
image_label_list = []
with open(filename, 'r') as f:
lines = f.readlines()
for line in lines:
# rstrip：用来去除结尾字符、空白符(包括\n、\r、\t、' '，即：换行、回车、制表符、空格)
content = line.rstrip().split(' ')
name = content[0]
labels = []
for value in content[1:]:
labels.append(int(value))
image_label_list.append((name, labels))
return image_label_list
def load_data(self, path, resize_height, resize_width, normalization):
'''
加载数据
:param path:
:param resize_height:
:param resize_width:
:param normalization: 是否归一化
:return:
'''
image = image_processing.read_image(path, resize_height, resize_width, normalization)
return image
def data_preproccess(self, data):
'''
数据预处理
:param data:
:return:
'''
data = self.toTensor(data)
return data
复制代码3.2 DataLoader产生批训练数据
'''
下面两种方式，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种迭代方法
for step, (batch_image, batch_label) in enumerate(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)
if step>=max_iterate:
break
# [3]第3种迭代方法
# for step in range(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)
复制代码