·请参考本系列目录：【mT5多语言翻译】之一——实战项目总览

[1] 加载数据集

  在上一篇实战博客中，我们介绍了如何下载和预处理数据集，并且介绍了如何将数据集内的文本进行分词然后保存为pt文件。

  接下来，我们讲解在模型训练时加载数据集的写法。

def build_dataset(config):def load_dataset(path, name):data = torch.cat(([torch.load(_) for _ in path]), dim=0)# 设置随机种子seed = 42torch.manual_seed(seed)# 生成一个随机排列random_permutation = torch.randperm(data.size(0))# 根据随机排列打乱张量data = data[random_permutation]for p in path:logger.info(f"{name}数据集地址========{p}")logger.info(f"数据集文本总数========{data.size(0)}")return datadev = load_dataset(config.dev_path, 'dev')train = load_dataset(config.train_path, 'train')return train, dev

  由于我们的项目有2个数据集（中韩翻译和中日翻译），因此config.dev_path和config.train_path是数组，内部分别是2个语料的pt文件地址：

'dev_path': [ '/home/20240321/@_Dataset_Processed/kor-zho/eval.pt','/home/20240321/@_Dataset_Processed/jpn-zho/eval.pt'],
'train_path': [ '/home/20240321/@_Dataset_Processed/kor-zho/train.pt','/home/20240321/@_Dataset_Processed/jpn-zho/train.pt']

  首先使用cat函数将两个数据集拼接起来，然后乱序打乱，最后返回即可。

  打乱的目的是让中韩翻译文本和中日翻译文本均匀分布，模型在训练时能够尽可能地均匀学习特征。

[2] 构造数据加载器

  数据加载器的作用是对上面加载的（1200万条）数据按照自定义的batch_size自动地进行批划分。代码如下：

class DatasetIterator(object):"""初始化时把数据加载到CPU迭代器每次加载batch_size个样本到GPU或CPU"""def __init__(self, batches, batch_size, device):if batch_size <= 0:raise ValueError("batch_size 必须大于 0。")self.batch_size = batch_sizeself.batches = batchesself.n_batches = len(batches) // batch_sizeself.has_residual = False  # 记录batch数量是否为整数if len(batches) % self.n_batches != 0:self.has_residual = Trueself.index = 0self.device = devicedef _to_tensor(self, datas):input_token_ids = torch.stack([data[0] for data in datas], dim=0).to(self.device)target_token_ids = torch.stack([data[1] for data in datas], dim=0).to(self.device)return input_token_ids, target_token_idsdef __next__(self):if self.has_residual and self.index == self.n_batches:batches = self.batches[self.index * self.batch_size: len(self.batches)]self.index += 1batches = self._to_tensor(batches)return batcheselif self.index >= self.n_batches:self.index = 0raise StopIterationelse:batches = self.batches[self.index * self.batch_size: (self.index + 1) * self.batch_size]self.index += 1batches = self._to_tensor(batches)return batchesdef __iter__(self):return selfdef __len__(self):if self.has_residual:return self.n_batches + 1else:return self.n_batches

  目前很多深度学习库都有成熟的数据加载器可以调用。这里选择手搓的原因是我觉得手搓的数据加载器更轻，可以非常直观的进行自定义，虽然代码少但是完全满足我们项目的功能需要。

【注】其实transformers库也提供了非常友好的代码训练api，可能都不用写多少代码，即可开始自定义任务、模型的训练。但是在使用过程中，还是遇到了一些问题：1）莫名其妙的多消耗显存，由于封装的功能太多，我很难定位到哪里出了问题；2）它提供的训练配置中有很多我不需要的功能，比如检查点保存、内部日志等等，同样也是集成的功能太多，但是我不需要。
————————————
在训练过程中，我想使项目代码尽可能的“轻量化”，并且直观、易于修改。第三方提供的训练方案都像黑盒一样，自定义起来比较麻烦。因此我还是选择了手写训练代码，并且代码并不多。

  DatasetIterator自定义了一些在迭代（for循环）时才会触发的函数。迭代时会先触发__next__函数，然后去调用_to_tensor函数将batch_size条数据加载到GPU上，之后送给模型训练。

  因此，也不同担心数据集太大，显存会承受不了的问题。

  项目中的1200万条数据先是加载到内存，然后每次只加载batch_size条数据加载到GPU上训练。

[3] 访问批数据

  定义好数据集加载函数和数据加载器之后，我们只需要写如下三行代码即可构造训练集和验证集的DataLoader：

train_data, dev_data = build_dataset(conf)
train_iter = build_iterator(train_data, conf)
dev_iter = build_iterator(dev_data, conf)

  最后我们写个循环去遍历DataLoader，即可访问批数据：

for i, (input_batch, label_batch) in enumerate(train_iter):passfor i, (input_batch, label_batch) in enumerate(dev_iter):pass

[4] 进行下一篇实战

  【mT5多语言翻译】之五——训练：中央日志、训练可视化、PEFT微调