内容

一、torch.utils.data.Dataset使用方法

1、NLP中数据量不是很大情况

2、图片和音频数据

二、data.IterableDataset的使用方法

做算法工程师工作以来，一直都是使用的pytorch框架，它提供了很方便的数据加载模块。很少接触到训练数据数据量巨大的情景，64G内存的机器采用Map加载方式(torch.utils.data.Dataset)一次性加载全部数据放在内存中也是能够做到的，训练效率也是比较高的，不需要每个batch实时读取数据。 但是当训练数据数据量大到一般的机器容纳不下的时候，采用Map加载方式就不行了，需要采用Iterable（迭代data.IterableDataset)，采用流式读取。一般而言，NLP中文本和CV中的图片以及语音领域的音频文件都有一套不同的读取方式。

一、torch.utils.data.Dataset使用方法

1、NLP中数据量不是很大情况

这个时候一般采用一次性把数据加载进内存的方式，NLP中分类任务、NER、翻译等一般都是使用transformer架构行模型，首先需要把字符映射到词典序号上。数据读取的时候先读取文本中所有的数据，做词典映射处理，得到模型输入向量。

class DataReader(Dataset):
    def __init__(self,tokenizer,filepath,max_len,task_type):
        self.tokenizer = tokenizer
        self.filepath = filepath
        self.max_len = max_len
        self.task_type = task_type
        self.dataList = self.datas_to_torachTensor()
        self.allLength = len(self.dataList)

    def convert_text2ids(self,text):
        text = text[0:self.max_len-2]
        inputs = self.tokenizer(text)

        input_ids = inputs['input_ids']
        attention_mask = inputs['attention_mask']
        paddings = [0] * (self.max_len - len(input_ids))
        input_ids += paddings
        attention_mask += paddings

        token_type_id = [0] * self.max_len

        return input_ids, attention_mask, token_type_id


    def datas_to_torachTensor(self):
         "一次性读取文本数据，通过相应的逻辑把文本转化为tensor张量等"
        convert_text2ids()
        return res

    def __getitem__(self, item):
        input_ids_a = self.dataList[item][0]
        attention_mask_a = self.dataList[item][1]
        token_type_id_a = self.dataList[item][2]

        input_ids_b = self.dataList[item][3]
        attention_mask_b = self.dataList[item][4]
        token_type_id_b = self.dataList[item][5]

        label = self.dataList[item][6]
        
        return input_ids_a, attention_mask_a, token_type_id_a, input_ids_b, attention_mask_b, token_type_id_b, label


    def __len__(self):
        return self.allLength

比较重要的就是重写__getitem__()、__len__()和__init__()方法。注意的是已经在__init__()把数据全部加载进了内存，所以训练过程中__getitem__()取值的时候非常快，训练的效率比较好，GPU利用率高。

2、图片和音频数据

这类数据每一条数据量都比较大，如果直接还是像上面一样加载进内存里面，内存压力就会比较大。可以直接在__init__()中把具体的语音或者图片数据的路径加载进内存，然后在__getitem__()根据路径实时读取对应数据。

class GCMasKCircleTimeDomain(Dataset):
    def __init__(self,data_dir):
        #音频数据的路径
        self.features_path = glob(data_dir+'/*bin')
        self.features_path.sort()

    def __getitem__(self, index):
        #实时加载音频特征(已经提前处理好做了本地持久化)
        feature = torch.load(self.features_path[index])[0]
        label = torch.load(self.features_path[index])[1]
        # # 如果没有持久化
        # feature = librosa.load(self.features_path[index],sr=16000)
        # import torchaudio
        # feature, sr = torchaudio.load(self.features_path[index])
        # ...
        # todo
        return feature, label

    def __len__(self):
        return len(self.features_path)

如上代码，首先在__init__()加载数据的路径、然后重写__getitem__()、__len__()，具体的在__getitem__()中实时读取数据。这种方式的优点就是可以完成大量数据的训练，就是每个batch的数据都是从硬盘上读取的，训练的时候GPU的使用率不能拉满。

二、data.IterableDataset的使用方法

有些情况下，不知道数据量的具体大小，就不能采用上述的map方式了，torch.utils.data.Dataset中的__len__()方法重写不了，没有具体长度。这个时候就需要使用Iteration方式了，可迭代的流式读取。

代码显示如下：

class MyIterableDataset(IterableDataset):

    def __init__(self, file_list):
        super(MyIterableDataset, self).__init__()
        self.file_list = file_list 

    def parse_file(self):
        for file in self.file_list: # 逐个文件读取
            print("读取文件：", file)
            with open(file, 'r') as file_obj:
                for line in file_obj: # 逐行读取
                	# 这里可以根据具体文件格式读取，进行逻辑处理
                    feature, label = self.extract(line)
                    yield (feature, label)
   
    def extract(self,line):
        #do something
        return

    def __iter__(self):
    	# 如果 batch_size = n 这会运行这个方法n次
        return self.parse_file()

注意这里的parse_file函数其实就是一个生成器了，可以在节约内存的同时实现遍历；

最后配合torch.utils.data.Dataloder()就可以构建模型训练流程的数据迭代器了，然后就可以愉快的炼丹了！

train_set = GCMasKCircleTimeDomain(data_dir='./data/time_domain_featres/train_final')
train_loader = DataLoader(train_set, batch_size=config.batch_size, collate_fn=pad_collate, shuffle=True)

train_set = MyIterableDataset(file_list = "./data/train.tsv")
train_loader = DataLoader(train_set, batch_size=config.batch_size, collate_fn=pad_collate, shuffle=True)