使用Paddle2.0构建一个经典的文本生成模型GPT-2

最新推荐文章于 2023-12-23 19:27:49 发布

原创

最新推荐文章于 2023-12-23 19:27:49 发布 · 2k 阅读

9 ·

CC 4.0 BY-SA版权

文章标签：

#paddlepaddle #python #深度学习

引入

GPT-2 是一种基于 transformer 的大型语言模型，是一个经典的文本生成模型
该模型可以生成连贯的文本段落，完成阅读理解、问答、机器翻译等多项不同的语言建模任务
本次就介绍一下如何使用Paddle2.0构建一个经典的文本生成模型GPT-2

环境设置

本文基于飞桨2.0-rc0版本

# 导入必要的包
import math
import paddle
import paddle.nn as nn

# 打印Paddle版本
print('paddle version: %s' % paddle.__version__)

paddle version: 2.0.0-rc0

模型架构

GPT-2是一个单向的语言模型，由前后的Embedding层和Transformer Encoder两部分组成
其中Transformer Encoder是由多层的Multi-Head-Self-Attention Layer+ MLP Layer堆叠而成的

构建 Multi-Head-Self-Attention Layer

首先从最基础的 Layer 开始构建
从名字就能看出这个层的特点：Multi-Head——多头、Self-Attention——自注意力

class Attention(nn.Layer):
    def __init__(self, 
                embedding_size, 
                num_attention_heads,
                attention_dropout,
                residual_dropout):
        super(Attention, self).__init__()
        
        self.num_attention_heads = num_attention_heads
        self.size_per_head = embedding_size // num_attention_heads
        self.embedding_size = embedding_size

        self.query_key_value = nn.Linear(embedding_size, embedding_size * 3)
        self.attn_drop = nn.Dropout(attention_dropout)
        self.resid_drop = nn.Dropout(residual_dropout)
        self.dense = nn.Linear(embedding_size, embedding_size)
    
    def split_heads(self, x):
        x = x.reshape([-1, self.seq_len, self.num_attention_heads, self.size_per_head])
        return x.transpose((0, 2, 1, 3))

    def forward(self, x, kv_cache=None):
        self.seq_len = x.shape[1]

        # 自注意力
        x = self.query_key_value(x)
        q, k, v = x.split(num_or_sections=3, axis=2)
        
        # 多头
        q = self.split_heads(q)
        k = self.split_heads(k)
        v = self.split_heads(v)
        
        # 缓存
        if kv_cache is not None:
            pk, pv = paddle.unstack(