【古诗生成AI实战】之四——模型包装器与模型的训练

  在上一篇博客中，我们已经利用任务加载器

task

成功地从数据集文件中加载了文本数据，并通过预处理器

processor

构建了词典和编码器。在这一过程中，我们还完成了词向量的提取。

  接下来的步骤涉及到定义模型、加载数据，并开始训练过程。

  为了确保项目代码能够快速切换到不同的模型，并且能够有效地支持

transformers

库中的预训练模型，我们不仅仅是定义模型那么简单。为此，我们采取了进一步的措施：在模型外面再套上一个额外的层，我称之为模型包装器

NNModelWrapper

。此外，为了提高配置的灵活性和可维护性，我们将所有的配置项（如批量大小、数据集地址、训练周期数、学习率等）抽取出来，统一放置在一个名为

WrapperConfig

的配置容器中。通过这种方式，我们就可以避免直接在代码中修改配置参数，而是通过更改配置文件来实现，从而使得整个项目更加模块化和易于管理。

  本章内容属于模型训练阶段，将分别介绍包装器配置

WrapperConfig

、模型包装器

NNModelWrapper

和模型

Model

。

[1] 包装器配置

WrapperConfig

  我们把配置全部放在yaml文件里，然后读取里面的配置，赋值给

WrapperConfig

类。定义如下：

classWrapperConfig(object):"""A configuration for a :class:`NNModelWrapper`."""def__init__(
            self,
            tokenizer,
            max_seq_len:int,
            vocab_num:int,
            word2vec_path:str,
            batch_size:int=1,
            epoch_num:int=1,
            learning_rate:float=0.001):
        self.tokenizer = tokenizer
        self.max_seq_len = max_seq_len
        self.batch_size = batch_size
        self.epoch_num = epoch_num
        self.learning_rate = learning_rate
        self.word2vec_path = word2vec_path
        self.vocab_num = vocab_num

WrapperConfig

类用于配置神经网络模型包装器（

NNModelWrapper

）。类的构造函数接受多个参数来初始化配置：

tokenizer: 分词器对象，用于文本处理或文本转换为模型可理解的格式。其实就是预处理器

processor

提供的

tokenizer

。

max_seq_len (int): 模型可以处理的最大序列长度。

vocab_num (int): 词汇表的大小。

word2vec_path (str):预训练的词向量模型的文件路径。即上文提取的词向量。

batch_size (int): 每个批次处理的数据样本数量。

epoch_num (int): 训练轮次。

learning_rate (float): 学习率。

[2] 模型包装器

NNModelWrapper

  模型包装器

NNModelWrapper

接受2个参数，一个是包装器配置

WrapperConfig

，另外一个是自定义模型

Model

。代码如下：

classNNModelWrapper:"""A wrapper around a Transformer-based language model."""def__init__(self, config: WrapperConfig, model):"""Create a new wrapper from the given config."""
        self.config = config
        self.model = model(self.config)defgenerate_dataset(self, data, labeled=True):"""Generate a dataset from the given examples."""
        features = self._convert_examples_to_features(data)

        feature_dict ={'input_ids': torch.tensor([f.input_ids for f in features], dtype=torch.long),'labels': torch.tensor([f.labels for f in features], dtype=torch.long),}ifnot labeled:del feature_dict['labels']return DictDataset(**feature_dict)def_convert_examples_to_features(self, examples)-> List[InputFeatures]:"""Convert a set of examples into a list of input features."""
        features =[]for(ex_index, example)in tqdm(enumerate(examples)):if ex_index %5000==0:
                logging.info("Writing example {}".format(ex_index))
            input_features = self.get_input_features(example)
            features.append(input_features)# logging.info(f"最终数据构造形式：{features[0]}")return features

    defget_input_features(self, example)-> InputFeatures:"""Convert the given example into a set of input features"""
        text = example.text
        input_ids = self.config.tokenizer(text)
        labels = np.copy(input_ids)
        labels[:-1]= input_ids[1:]assertlen(input_ids)== self.config.max_seq_len

        return InputFeatures(input_ids=input_ids, attention_mask=None, token_type_ids=None, labels=labels)

NNModelWrapper

类是围绕一个神经网络语言模型的封装器，提供了模型的初始化和数据处理的方法。

  · 类初始化 (init):

config

: 接收一个

WrapperConfig

类的实例，包含模型的配置信息。

model

: 接收一个模型构造函数，该函数使用配置信息来初始化模型。

  · 生成数据集 (generate_dataset):从给定的数据样本中生成一个数据集。首先把数据样本转换为特征（通过

_convert_examples_to_features

方法），然后根据这些特征创建一个

DictDataset

对象。如果数据未标记（

labeled=False

），则从特征字典中删除

labels

键。

  · 转换样本为特征 (_convert_examples_to_features):这是个私有方法，把数据样本转换为模型可以理解的输入特征。对于每个样本，使用

get_input_features

方法来生成输入特征。使用

tqdm

显示处理进度，并利用

logging

记录处理信息。

  · 获取输入特征 (get_input_features):此方法将单个数据样本转换为输入特征。首先获取文本内容，然后使用配置中的分词器（

tokenizer

）将文本转换为

input_ids

。标签（

labels

）是

input_ids

的一个变体，其中每个元素都向右移动一个位置。用断言确保

input_ids

的长度与配置中的

max_seq_len

相等。

[3] 模型

Model

  模型包装器

NNModelWrapper

里面的第二个参数

Model

才是我们真正的模型。

  在古诗生成AI任务中，RNN是比较适配任务的模型，我们定义的RNN模型如下：

classModel(nn.Module):def__init__(self, config):super(Model, self).__init__()

        V = config.vocab_num  # vocab_num
        E =300# embed_dim
        H =256# hidden_size

        embedding_pretrained = torch.tensor(np.load(config.word2vec_path)["embeddings"].astype('float32'))
        self.embedding = nn.Embedding.from_pretrained(embedding_pretrained, freeze=False)
        self.lstm = nn.LSTM(E, H,1, bidirectional=False, batch_first=True, dropout=0.1)
        self.fc = nn.Linear(H, V)
        self.loss = nn.CrossEntropyLoss()defforward(self, input_ids, labels=None):
        embed = self.embedding(input_ids)# [batch_size, seq_len, embed_dim]
        out, _ = self.lstm(embed)# [batch_size, seq_len, hidden_size]
        logit = self.fc(out)# [batch_size, seq_len, vocab_num]if labels isnotNone:
            loss = self.loss(logit.view(-1, logit.shape[-1]), labels.view(-1))return loss, logit
        else:return logit[None,:]

  在我们的模型中，特别值得一提的是嵌入层（embedding layer）。在初始化这一层时，我们使用的是之前提取出的词向量。这种做法有助于模型更好地理解和处理文本数据。

  此次我们定义的模型是一个基于

RNN

的结构，它包括三个主要部分：

embedding

层、

lstm

层和

fc

（全连接）层。

  在模型的前向传播（

forward

）过程中，输入

input_ids

的形状为

[batch_size, seq_len]

，即每个批次有多少文本，每个文本的序列长度是多少。输入数据首先通过嵌入层处理，输出的

embed

的形状为

[batch_size, seq_len, embed_dim]

，即每个单词都被转换成了对应的嵌入向量。接着，数据通过一个单层的

lstm

网络，得到输出

out

，最后经过全连接层

fc

，得到最终的概率分布

logit

。

  这个概率分布

logit

的含义是：对于每个批次中的文本，每个文本在序列的每个位置上，都有vocab_num个可能的词可以填入，而

logit

中存储的正是这些词的概率。为了生成文本，我们提取每个位置上概率最高的词的索引，然后根据这些索引在词典中查找对应的词。这就是我们通过模型运行文本生成得到的结果。

[4] 训练

  所有的工作都准备好了，下面我们正式开始模型的训练。

  对于神经网络的训练、验证、测试、优化等等操作，我采用了

transformers

的

Trainer

极大的简化了项目操作。

  第一步，加载yaml配置文件，读取所有配置项：

withopen('config.yaml','r', encoding='utf-8')as f:
        conf = yaml.load(f.read(),Loader=yaml.FullLoader)
        conf_train = conf['train']
        conf_sys = conf['sys']

  第二步，初始化任务加载器，加载数据集：

    Task = TASKS[conf_train['task_name']]()

    data = Task.get_train_examples(conf_train['dataset_url'])
    index =int(len(data)* conf_train['rate'])
    train_data, dev_data = data[:index], data[index:]

  第三步，初始化数据预处理器，并向外提供

tokenizer

：

    Processor = PROCESSORS[conf_train['task_name']](data, conf_train['max_seq_len'], conf_train['vocab_path'])
    tokenizer =lambda text: Processor.tokenizer(text)

  第四步，初始化模型包装配置：

    wrapper_config = WrapperConfig(
        tokenizer=tokenizer,
        max_seq_len=conf_train['max_seq_len'],
        batch_size=conf_train['batch_size'],
        epoch_num=conf_train['epoch_num'],
        learning_rate=conf_train['learning_rate'],
        word2vec_path=conf_train['word2vec_path'],
        vocab_num=len(Processor.vocab))

  第五步，加载模型，初始化模型包装器：

    x = import_module(f'main.model.{conf_train["model_name"]}')
    wrapper = NNModelWrapper(wrapper_config, x.Model)print(f'模型有 {sum(p.numel()for p in wrapper.model.parameters()if p.requires_grad):,} 个训练参数')

  第六步，使用模型包装器生成数据集向量：

train_dataset = wrapper.generate_dataset(train_data)
    val_dataset = wrapper.generate_dataset(dev_data)

  第七步，创建训练器：

# 构建trainerdefcreate_trainer(wrapper, train_dataset, val_dataset):# 模型
    model = wrapper.model

    args = TrainingArguments('./checkpoints',# 模型保存的输出目录
        save_strategy=IntervalStrategy.STEPS,# 模型保存策略
        save_steps=50,# 每n步保存一次模型  1步表示一个batch训练结束
        evaluation_strategy=IntervalStrategy.STEPS,
        eval_steps=50,
        overwrite_output_dir=True,# 设置overwrite_output_dir参数为True，表示覆盖输出目录中已有的模型文件
        logging_dir='./logs',# 可视化数据文件存储地址
        log_level="warning",
        logging_steps=50,# 每n步保存一次评价指标  1步表示一个batch训练结束 | 还控制着控制台的打印频率 每n步打印一下评价指标 | n过大时，只会保存最后一次的评价指标
        disable_tqdm=True,# 是否不显示数据训练进度条
        learning_rate=wrapper.config.learning_rate,
        per_device_train_batch_size=wrapper.config.batch_size,
        per_device_eval_batch_size=wrapper.config.batch_size,
        num_train_epochs=wrapper.config.epoch_num,
        dataloader_num_workers=2,# 数据加载的子进程数
        weight_decay=0.01,
        save_total_limit=2,
        load_best_model_at_end=True)# 早停设置
    early_stopping = EarlyStoppingCallback(
        early_stopping_patience=3,# 如果8次验证集性能没有提升，则停止训练
        early_stopping_threshold=0,# 验证集的性能提高不到0时也停止训练)

    trainer = Trainer(
        model,
        args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        callbacks=[early_stopping],# 添加EarlyStoppingCallback回调函数)return trainer
    
trainer = create_trainer(wrapper, train_dataset, val_dataset)

  第八步，开始训练并设置保存模型：

    trainer.train()
    trainer.save_model(conf_train['model_save_dir']+ conf_train['task_name']+'/'+ conf_train['model_name'])

  训练的整体代码如下：

# 构建trainerdefcreate_trainer(wrapper, train_dataset, val_dataset):# 模型
    model = wrapper.model

    args = TrainingArguments('./checkpoints',# 模型保存的输出目录
        save_strategy=IntervalStrategy.STEPS,# 模型保存策略
        save_steps=50,# 每n步保存一次模型  1步表示一个batch训练结束
        evaluation_strategy=IntervalStrategy.STEPS,
        eval_steps=50,
        overwrite_output_dir=True,# 设置overwrite_output_dir参数为True，表示覆盖输出目录中已有的模型文件
        logging_dir='./logs',# 可视化数据文件存储地址
        log_level="warning",
        logging_steps=50,# 每n步保存一次评价指标  1步表示一个batch训练结束 | 还控制着控制台的打印频率 每n步打印一下评价指标 | n过大时，只会保存最后一次的评价指标
        disable_tqdm=True,# 是否不显示数据训练进度条
        learning_rate=wrapper.config.learning_rate,
        per_device_train_batch_size=wrapper.config.batch_size,
        per_device_eval_batch_size=wrapper.config.batch_size,
        num_train_epochs=wrapper.config.epoch_num,
        dataloader_num_workers=2,# 数据加载的子进程数
        weight_decay=0.01,
        save_total_limit=2,
        load_best_model_at_end=True)# 早停设置
    early_stopping = EarlyStoppingCallback(
        early_stopping_patience=3,# 如果8次验证集性能没有提升，则停止训练
        early_stopping_threshold=0,# 验证集的性能提高不到0时也停止训练)

    trainer = Trainer(
        model,
        args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        callbacks=[early_stopping],# 添加EarlyStoppingCallback回调函数)return trainer

defmain():# ### @通用配置# ##withopen('config.yaml','r', encoding='utf-8')as f:
        conf = yaml.load(f.read(),Loader=yaml.FullLoader)
        conf_train = conf['train']
        conf_sys = conf['sys']# 系统设置初始化
    System(conf_sys).init_system()# 初始化任务加载器
    Task = TASKS[conf_train['task_name']]()

    data = Task.get_train_examples(conf_train['dataset_url'])
    index =int(len(data)* conf_train['rate'])
    train_data, dev_data = data[:index], data[index:]# 初始化数据预处理器
    Processor = PROCESSORS[conf_train['task_name']](data, conf_train['max_seq_len'], conf_train['vocab_path'])
    tokenizer =lambda text: Processor.tokenizer(text)# 初始化模型包装配置
    wrapper_config = WrapperConfig(
        tokenizer=tokenizer,
        max_seq_len=conf_train['max_seq_len'],
        batch_size=conf_train['batch_size'],
        epoch_num=conf_train['epoch_num'],
        learning_rate=conf_train['learning_rate'],
        word2vec_path=conf_train['word2vec_path'],
        vocab_num=len(Processor.vocab))

    x = import_module(f'main.model.{conf_train["model_name"]}')
    wrapper = NNModelWrapper(wrapper_config, x.Model)print(f'模型有 {sum(p.numel()for p in wrapper.model.parameters()if p.requires_grad):,} 个训练参数')# 生成数据集
    train_dataset = wrapper.generate_dataset(train_data)
    val_dataset = wrapper.generate_dataset(dev_data)# 训练与保存
    trainer = create_trainer(wrapper, train_dataset, val_dataset)
    trainer.train()
    trainer.save_model(conf_train['model_save_dir']+ conf_train['task_name']+'/'+ conf_train['model_name'])if __name__ =='__main__':
    main()

  运行之后，看到下面输出代表项目成功运行：

[5] 进行下一篇实战

  【古诗生成AI实战】之五——加载模型进行古诗生成