在上一篇博客中,我们已经利用任务加载器task
成功地从数据集文件中加载了文本数据,并通过预处理器processor
构建了词典和编码器。在这一过程中,我们还完成了词向量的提取。
接下来的步骤涉及到定义模型、加载数据,并开始训练过程。
为了确保项目代码能够快速切换到不同的模型,并且能够有效地支持transformers
库中的预训练模型,我们不仅仅是定义模型那么简单。为此,我们采取了进一步的措施:在模型外面再套上一个额外的层,我称之为模型包装器NNModelWrapper
。此外,为了提高配置的灵活性和可维护性,我们将所有的配置项(如批量大小、数据集地址、训练周期数、学习率等)抽取出来,统一放置在一个名为WrapperConfig
的配置容器中。通过这种方式,我们就可以避免直接在代码中修改配置参数,而是通过更改配置文件来实现,从而使得整个项目更加模块化和易于管理。
本章内容属于模型训练阶段,将分别介绍包装器配置WrapperConfig
、模型包装器NNModelWrapper
和模型Model
。
[1] 包装器配置WrapperConfig
我们把配置全部放在yaml文件里,然后读取里面的配置,赋值给WrapperConfig
类。定义如下:
class WrapperConfig(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
。代码如下:
class NNModelWrapper:
"""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)
def generate_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),
}
if not 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
def get_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:]
assert len(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模型如下:
class Model(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()
def forward(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 is not None:
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配置文件,读取所有配置项:
with open('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)
第七步,创建训练器:
# 构建trainer
def create_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'])
训练的整体代码如下:
# 构建trainer
def create_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
def main():
# ##
# @通用配置
# ##
with open('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实战】之五——加载模型进行古诗生成