gh_mirrors/py/pytorch-book项目案例：用Transformer实现文本生成
项目地址: https://gitcode.com/gh_mirrors/py/pytorch-book 你是否曾想过让计算机像古代诗人一样创作优美的诗句？或者输入几个关键词就能生成完整的文章？在深度学习框架PyTorch中，这一切都变得触手可及。本文将带你深入了解如何使用Transformer模型实现文本生成，通过gh_mirrors/py/pytorch-book项目中的Chapter11案例，你将学会从数据准备到模型训练，再到生成诗词的完整流程。读完本文，你将能够：掌握Transformer模型的基本原理、理解文本生成的实现过程、运行并修改项目代码生成自己的文本内容。
Transformer模型简介
Transformer是一种基于自注意力机制（Self-Attention）的神经网络模型，由Vaswani等人在2017年提出。与传统的循环神经网络（RNN）相比，Transformer具有并行计算能力强、长距离依赖捕捉效果好等优点，因此在自然语言处理任务中得到了广泛应用。
在gh_mirrors/py/pytorch-book项目的Chapter11中，Transformer模型被用于诗歌生成任务。该模型主要由以下几个部分组成：
- 词嵌入（Token Embedding）：将输入的文本转换为向量表示。
- 位置编码（Positional Encoding）：为每个位置添加位置信息，使模型能够捕捉序列的顺序特征。
- 编码器（Encoder）：由多个编码器层组成，每个编码器层包含多头自注意力机制和前馈神经网络。
项目结构与文件解析
gh_mirrors/py/pytorch-book项目的Chapter11目录下包含了实现Transformer文本生成的所有代码文件，主要包括：
- main.py：项目的主入口文件，包含训练和生成文本的主要逻辑。
- model.py：定义了Transformer模型的结构，包括词嵌入、位置编码和编码器等。
- data.py：数据预处理模块，负责读取原始文本数据、构建词汇表和生成训练数据。
- utils.py：工具函数模块，包含可视化等辅助功能。
下面我们将重点解析其中的关键文件和代码。
model.py：Transformer模型定义
在model.py文件中，定义了PoetryModel类，该类继承自nn.Module，是实现文本生成的核心。
首先，词嵌入模块将输入的词转换为向量表示：
class TokenEmbedding(nn.Module):
def __init__(self, vocab_size, emb_size):
super().__init__()
self.embedding = nn.Embedding(vocab_size, emb_size)
self.emb_size = emb_size
def forward(self, tokens):
return self.embedding(tokens.long()) * math.sqrt(self.emb_size)
位置编码模块为每个位置添加位置信息：
class PositionalEncoding(nn.Module):
def __init__(self, emb_size, dropout, maxlen=200):
super().__init__()
den = torch.exp(- torch.arange(0, emb_size, 2).float() * math.log(100) / emb_size)
pos = torch.arange(0, maxlen).float().reshape(maxlen, 1)
pos_embedding = torch.zeros((maxlen, emb_size))
pos_embedding[:, 0::2] = torch.sin(pos * den)
pos_embedding[:, 1::2] = torch.cos(pos * den)
pos_embedding = pos_embedding.unsqueeze(-2)
self.dropout = nn.Dropout(dropout)
self.register_buffer('pos_embedding', pos_embedding)
def forward(self, token_embedding):
return self.dropout(token_embedding + self.pos_embedding[:token_embedding.size(0),:])
PoetryModel类将词嵌入、位置编码和编码器组合在一起：
class PoetryModel(nn.Module):
def __init__(self, vocab_size, num_encoder_layers=4, emb_size=512, dim_feedforward=1024, dropout=0.1):
super().__init__()
self.src_tok_emb = TokenEmbedding(vocab_size, emb_size)
self.positional_encoding = PositionalEncoding(emb_size, dropout=dropout)
encoder_layer = TransformerEncoderLayer(d_model=emb_size, nhead=8, dim_feedforward=dim_feedforward)
self.transformer_encoder = TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
self.generator = nn.Linear(emb_size, vocab_size)
def forward(self, src, src_mask, src_padding_mask):
src_emb = self.positional_encoding(self.src_tok_emb(src))
memory = self.transformer_encoder(src_emb, src_mask, src_padding_mask)
logit = self.generator(memory)
return memory, logit
main.py：训练与生成逻辑
在main.py文件中，定义了训练和生成文本的函数。
训练函数train()主要包括数据加载、模型初始化、损失函数定义和优化器设置等步骤：
def train(**kwargs):
for k, v in kwargs.items():
setattr(opt, k, v)
device = t.device('cuda') if opt.use_gpu else t.device('cpu')
vis = Visualizer(env = opt.env)
# 获取数据
data, word2ix, ix2word = get_data(opt)
data = t.from_numpy(data)
dataloader = DataLoader(data, batch_size=opt.batch_size, shuffle=True, num_workers=1)
# 模型定义
model = PoetryModel(len(word2ix))
optimizer = t.optim.Adam(model.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9)
criterion = nn.CrossEntropyLoss(ignore_index=len(word2ix)-1)
if opt.model_path:
model.load_state_dict(t.load(opt.model_path))
model.to(device)
# 训练循环
for epoch in range(opt.epoch):
loss_meter.reset()
for ii, data_ in tqdm.tqdm(enumerate(dataloader)):
# 训练步骤
data_ = data_.long().transpose(1, 0).contiguous()
data_ = data_.to(device)
optimizer.zero_grad()
input_, target = data_[:-1, :], data_[1:, :]
src_mask = generate_square_subsequent_mask(input_.shape[0])
src_pad_mask = input_ == len(word2ix) - 1
src_pad_mask = src_pad_mask.permute(1,0).contiguous()
memory, logit = model(input_, src_mask.to(device), src_pad_mask.to(device))
mask = target != word2ix['</s>']
target = target[mask]
logit = logit.flatten(0, 1)[mask.view(-1)]
loss = criterion(logit, target)
loss.backward()
optimizer.step()
loss_meter.add(loss.item())
生成函数gen()用于根据输入的起始文本生成完整的诗句：
def gen(**kwargs):
"""
给定几个词，根据这几个词接着生成一首完成的诗词
"""
for k, v in kwargs.items():
setattr(opt, k, v)
device = t.device('cuda') if opt.use_gpu else t.device('cpu')
data, word2ix, ix2word = get_data(opt)
model = PoetryModel(len(word2ix))
model.load_state_dict(t.load(opt.model_path))
model.to(device)
model.eval()
src = [word2ix[word] for word in opt.start_words]
res = src = [word2ix['<START>']] + src
max_len = 100
for _ in range(max_len):
src = t.tensor(res).to(device)[:, None]
src_mask = generate_square_subsequent_mask(src.shape[0])
src_pad_mask = src == len(word2ix) - 1
src_pad_mask = src_pad_mask.permute(1, 0).contiguous()
memory, logits = model(src, src_mask.cuda(), src_pad_mask.cuda())
next_word = logits[-1, 0].argmax().item()
if next_word == word2ix['<EOP>']:
break
res.append(next_word)
res = [ix2word[_] for _ in res[1:]]
print(''.join(res))
此外，main.py中还定义了gen_acrostic()函数，用于生成藏头诗。
data.py：数据预处理
data.py文件中的get_data()函数负责读取原始文本数据、构建词汇表和生成训练数据：
def get_data(opt):
if os.path.exists(opt.pickle_path): # 预处理好的二进制文件存在
data = np.load(opt.pickle_path, allow_pickle=True)
data, word2ix, ix2word = data['data'], data['word2ix'].item(), data['ix2word'].item()
return data, word2ix, ix2word
# 如果没有处理好的二进制文件，则处理原始的json文件
data = _parseRawData(opt.author, opt.constrain, opt.data_path, opt.category)
words = {_word for _sentence in data for _word in _sentence}
word2ix = {_word: _ix for _ix, _word in enumerate(words)}
word2ix['<EOP>'] = len(word2ix) # 终止标识符
word2ix['<START>'] = len(word2ix) # 起始标识符
word2ix['</s>'] = len(word2ix) # 空格
ix2word = {_ix: _word for _word, _ix in list(word2ix.items())}
# 为每首诗歌加上起始符和终止符
for i in range(len(data)):
data[i] = ["<START>"] + list(data[i]) + ["<EOP>"]
# 将每首诗歌保存的内容由‘字’变成‘数’
new_data = [[word2ix[_word] for _word in _sentence] for _sentence in data]
# 诗歌长度不够opt.maxlen的在后面补空格，超过的opt.maxlen，删除末尾的字符
pad_data = pad_sequences(new_data, maxlen=opt.maxlen, padding='post', truncating='post', value=len(word2ix) - 1)
# 保存成二进制文件
np.savez_compressed(opt.pickle_path, data=pad_data, word2ix=word2ix, ix2word=ix2word)
return pad_data, word2ix, ix2word
运行与使用方法
环境准备
首先，需要安装项目所需的依赖库，可以通过Chapter11目录下的requirements.txt文件安装：
pip install -r Chapter11/requirements.txt
训练模型
运行main.py文件中的train()函数可以开始训练模型：
python Chapter11/main.py train --model_path checkpoints/tang --epoch 200 --batch_size 128
其中，--model_path指定模型保存路径，--epoch指定训练轮数，--batch_size指定批次大小。
生成文本
训练完成后，可以使用gen()函数生成文本。例如，生成以“深度学习”开头的诗句：
python Chapter11/main.py gen --model_path checkpoints/tang_200.pth --start_words 深度学习
生成的诗句可能如下：
深度学习，智识开启新境界。算法迭代，模型优化无止境。数据洪流，洞察规律见真章。人工智能，赋能未来创可能。
如果要生成藏头诗，可以使用gen_acrostic()函数：
python Chapter11/main.py gen_acrostic --model_path checkpoints/tang_200.pth --start_words 深度学习
生成的藏头诗可能如下：
深院静，小庭空，断续寒砧断续风。 度月影才敛，绕竹光复流。 学道无成鬓已霜，钓鱼犹得返家乡。 习隐深山里，烟霞是往还。
模型优化与改进
为了提高文本生成的质量，可以从以下几个方面对模型进行优化和改进：
- 增加训练数据：使用更多的文本数据进行训练，可以提高模型的泛化能力。
- 调整模型参数：如增加编码器层数、调整隐藏层维度、改变注意力头数等。
- 使用预训练模型：可以使用BERT等预训练模型作为词嵌入层，提高模型的初始化质量。
- 改进解码策略：除了贪婪搜索外，可以使用束搜索（Beam Search）等解码策略，生成更优的文本。
总结与展望
本文详细介绍了gh_mirrors/py/pytorch-book项目中使用Transformer实现文本生成的案例，包括模型结构、项目文件解析、运行方法和模型优化等方面。通过这个案例，我们可以看到Transformer模型在文本生成任务中的强大能力。
未来，我们可以进一步探索Transformer模型的变体，如GPT（Generative Pre-trained Transformer）等，并将其应用于更多的自然语言处理任务，如机器翻译、文本摘要和对话系统等。
希望本文能够帮助你更好地理解和应用Transformer模型进行文本生成。如果你对项目有任何疑问或建议，欢迎在项目的GitHub仓库中提出。
项目地址
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
