#

可配置变压器组件

9import copy
10
11import torch.nn as nn
12
13from labml.configs import BaseConfigs, option, calculate, aggregate
14from labml_helpers.module import Module
15from .feed_forward import FeedForward
16from .mha import MultiHeadAttention
17from .models import EmbeddingsWithPositionalEncoding, EmbeddingsWithLearnedPositionalEncoding, TransformerLayer, \
18    Encoder, Decoder, Generator, EncoderDecoder

#

FFN 配置

创建在中定义的位置前馈网络feed_forward.py 。

21class FeedForwardConfigs(BaseConfigs):

#

位置前馈层

31    ffn: FeedForward

#

嵌入中的要素数量

33    d_model: int

#

隐藏图层中的要素数量

35    d_ff: int = 2048

#

辍学概率

37    dropout: float = 0.1

#

在位置前馈层激活

39    activation: nn.Module = 'ReLU'

#

是否应对 FFN 层进行门控

41    is_gated: bool = False

#

第一个完全连接的层是否应该有可学习的偏差

43    bias1: bool = True

#

第二个全连接层是否应该有可学习的偏差

45    bias2: bool = True

#

栅极的全连接层是否应具有可学习的偏差

47    bias_gate: bool = False

#

预定义的 GLU 变体

49    glu_variant: str = 'none'

#

激活 ReLU

$max (0, x)$

52@option(FeedForwardConfigs.activation, 'ReLU')
53def _ffn_activation_relu():

#

59    return nn.ReLU()

#

GELU 激活

$x Φ (x)$ 在哪里 $Φ (x) = P (X \leq x), X \sim N (0, 1)$

它是在论文中介绍的 “高斯误差线性单位”。

62@option(FeedForwardConfigs.activation, 'GELU')
63def _ffn_activation_gelu():

#

71    return nn.GELU()

#

初始化前馈网络

74@option(FeedForwardConfigs.ffn, 'default')
75def _feed_forward(c: FeedForwardConfigs):

#

79    return FeedForward(c.d_model, c.d_ff,
80                       dropout=c.dropout,
81                       activation=c.activation,
82                       is_gated=c.is_gated,
83                       bias1=c.bias1,
84                       bias2=c.bias2,
85                       bias_gate=c.bias_gate)

#

GLU 变体

这些是用于FFN的封闭隐藏层的变体，如纸质 GLU变体改进变压器中所述。我们省略了本文中指定的偏差术语。

#

带门控线性单元的 FFN

$FF N_{G LU} (x) (x, W_{1}, V, W_{2}) = (σ (x W_{1}) \otimes x V) W_{2}$

95aggregate(FeedForwardConfigs.glu_variant, 'GLU',
96          (FeedForwardConfigs.is_gated, True),
97          (FeedForwardConfigs.bias1, False),
98          (FeedForwardConfigs.bias2, False),
99          (FeedForwardConfigs.bias_gate, False),
100          (FeedForwardConfigs.activation, nn.Sigmoid()))

#

带双线性隐藏层的 FFN

$FF N_{B i l in e a r} (x) (x, W_{1}, V, W_{2}) = (x W_{1} \otimes x V) W_{2}$

105aggregate(FeedForwardConfigs.glu_variant, 'Bilinear',
106          (FeedForwardConfigs.is_gated, True),
107          (FeedForwardConfigs.bias1, False),
108          (FeedForwardConfigs.bias2, False),
109          (FeedForwardConfigs.bias_gate, False),
110          (FeedForwardConfigs.activation, nn.Identity()))

#

带有 ReLU 门的 FFN

$FF N_{R e G LU} (x) (x, W_{1}, V, W_{2}) = (max (0, x W_{1}) \otimes x V) W_{2}$

115aggregate(FeedForwardConfigs.glu_variant, 'ReGLU',
116          (FeedForwardConfigs.is_gated, True),
117          (FeedForwardConfigs.bias1, False),
118          (FeedForwardConfigs.bias2, False),
119          (FeedForwardConfigs.bias_gate, False),
120          (FeedForwardConfigs.activation, nn.ReLU()))

#

带有 GELU 门的 FFN

$FF N_{GEG LU} (x) (x, W_{1}, V, W_{2}) = (GELU (x W_{1}) \otimes x V) W_{2}$

125aggregate(FeedForwardConfigs.glu_variant, 'GEGLU',
126          (FeedForwardConfigs.is_gated, True),
127          (FeedForwardConfigs.bias1, False),
128          (FeedForwardConfigs.bias2, False),
129          (FeedForwardConfigs.bias_gate, False),
130          (FeedForwardConfigs.activation, nn.GELU()))

#

FFN 带 Swish gate

$FF N_{Sw i G LU} (x) (x, W_{1}, V, W_{2}) = (Swish_{1} (x W_{1}) \otimes x V) W_{2}$ 在哪里 $Swish_{β} (x) = x σ (β x)$

136aggregate(FeedForwardConfigs.glu_variant, 'SwiGLU',
137          (FeedForwardConfigs.is_gated, True),
138          (FeedForwardConfigs.bias1, False),
139          (FeedForwardConfigs.bias2, False),
140          (FeedForwardConfigs.bias_gate, False),
141          (FeedForwardConfigs.activation, nn.SiLU()))

#

变压器配置

这定义了变压器的配置。配置是使用选项函数计算的。这些是延迟加载的，因此只计算必要的模块。

144class TransformerConfigs(BaseConfigs):

#

注意头数量

156    n_heads: int = 8

#

变压器嵌入尺寸

158    d_model: int = 512

#

层数

160    n_layers: int = 6

#

辍学概率

162    dropout: float = 0.1

#

源词汇表中的标记数（用于令牌嵌入）

164    n_src_vocab: int

#

目标词汇表中的标记数（用于生成预测的对数）

166    n_tgt_vocab: int

#

编码器自我注意

169    encoder_attn: MultiHeadAttention = 'mha'

#

解码器自我注意

171    decoder_attn: MultiHeadAttention = 'mha'

#

解码器内存注意事项

173    decoder_mem_attn: MultiHeadAttention = 'mha'

#

可配置的前馈层

176    ffn: FeedForwardConfigs

#

编码器层

179    encoder_layer: TransformerLayer = 'default'

#

解码器层

181    decoder_layer: TransformerLayer = 'default'

#

由多个编码器层组成的编码器

184    encoder: Encoder = 'default'

#

由多个解码器层组成的编码器

186    decoder: Decoder = 'default'

#

源的嵌入层

189    src_embed: Module = 'fixed_pos'

#

目标嵌入层（用于解码器）

191    tgt_embed: Module = 'fixed_pos'

#

用于预测的 Logit 生成器

194    generator: Generator = 'default'

#

编码器-解码器

197    encoder_decoder: EncoderDecoder

#

多头注意

201def _mha(c: TransformerConfigs):
202    return MultiHeadAttention(c.n_heads, c.d_model, dropout_prob=c.dropout)
203
204
205calculate(TransformerConfigs.encoder_attn, 'mha', _mha)
206calculate(TransformerConfigs.decoder_attn, 'mha', _mha)
207calculate(TransformerConfigs.decoder_mem_attn, 'mha', _mha)

#

相对多头注意力

211def _relative_mha(c: TransformerConfigs):
212    from labml_nn.transformers.xl.relative_mha import RelativeMultiHeadAttention
213    return RelativeMultiHeadAttention(c.n_heads, c.d_model)
214
215
216calculate(TransformerConfigs.encoder_attn, 'relative', _relative_mha)
217calculate(TransformerConfigs.decoder_attn, 'relative', _relative_mha)
218calculate(TransformerConfigs.decoder_mem_attn, 'relative', _relative_mha)

#

创建前馈层配置

221@option(TransformerConfigs.ffn, 'default')
222def _feed_forward(c: TransformerConfigs):

#

226    conf = FeedForwardConfigs()
227    conf.set_default(FeedForwardConfigs.d_model, func=lambda: c.d_model)
228    conf.set_default(FeedForwardConfigs.dropout, func=lambda: c.dropout)
229    return conf

#

编码器层

232@option(TransformerConfigs.encoder_layer, 'default')
233def _encoder_layer(c: TransformerConfigs):

#

237    return TransformerLayer(d_model=c.d_model, self_attn=c.encoder_attn,
238                            src_attn=None, feed_forward=copy.deepcopy(c.ffn.ffn),
239                            dropout_prob=c.dropout)

#

解码器层

242@option(TransformerConfigs.decoder_layer, 'default')
243def _decoder_layer(c: TransformerConfigs):

#

247    return TransformerLayer(d_model=c.d_model, self_attn=c.decoder_attn,
248                            src_attn=c.decoder_mem_attn, feed_forward=copy.deepcopy(c.ffn.ffn),
249                            dropout_prob=c.dropout)

#

编码器

252@option(TransformerConfigs.encoder, 'default')
253def _encoder(c: TransformerConfigs):

#

257    return Encoder(c.encoder_layer, c.n_layers)

#

解码器

260@option(TransformerConfigs.decoder, 'default')
261def _decoder(c: TransformerConfigs):

#

265    return Decoder(c.decoder_layer, c.n_layers)

#

Logit 生成器

268@option(TransformerConfigs.generator, 'default')
269def _generator(c: TransformerConfigs):

#

273    return Generator(c.n_tgt_vocab, c.d_model)

#

固定位置嵌入

使用固定位置编码进行源嵌入

277@option(TransformerConfigs.src_embed, 'fixed_pos')
278def _src_embed_with_positional(c: TransformerConfigs):

#

282    return EmbeddingsWithPositionalEncoding(c.d_model, c.n_src_vocab)

#

使用固定位置编码进行目标嵌入

285@option(TransformerConfigs.tgt_embed, 'fixed_pos')
286def _tgt_embed_with_positional(c: TransformerConfigs):

#

290    return EmbeddingsWithPositionalEncoding(c.d_model, c.n_tgt_vocab)

#

学习过的位置嵌入

使用学习的位置编码进行源嵌入

294@option(TransformerConfigs.src_embed, 'learned_pos')
295def _src_embed_with_learned_positional(c: TransformerConfigs):

#

299    return EmbeddingsWithLearnedPositionalEncoding(c.d_model, c.n_src_vocab)

#

使用学习的位置编码进行目标嵌入

302@option(TransformerConfigs.tgt_embed, 'learned_pos')
303def _tgt_embed_with_learned_positional(c: TransformerConfigs):

#

307    return EmbeddingsWithLearnedPositionalEncoding(c.d_model, c.n_tgt_vocab)

#

没有位置嵌入

不带位置编码的源代码嵌入

311@option(TransformerConfigs.src_embed, 'no_pos')
312def _src_embed_without_positional(c: TransformerConfigs):

#

316    return nn.Embedding(c.n_src_vocab, c.d_model)

#

319@option(TransformerConfigs.tgt_embed, 'no_pos')
320def _tgt_embed_without_positional(c: TransformerConfigs):
321    return nn.Embedding(c.n_tgt_vocab, c.d_model)
322
323
324@option(TransformerConfigs.encoder_decoder, 'default')
325def _encoder_decoder(c: TransformerConfigs):
326    return EncoderDecoder(c.encoder, c.decoder, c.src_embed, c.tgt_embed, c.generator)