#

模糊平铺激活实验

在这里，我们训练一个在前馈网络中使用模糊平铺激活的变压器。我们用它来制作语言模型，然后在 Tiny Shakespeare 数据集上训练它进行演示。

但是，这可能不是自由贸易协定的理想任务，我们认为自由贸易协定更适合使用连续变量对数据进行建模。

22import copy
23
24import torch
25import torch.nn as nn
26
27from labml import experiment
28from labml.configs import option
29from labml_helpers.module import Module
30from labml_nn.activations.fta import FTA
31from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
32from labml_nn.transformers import MultiHeadAttention, TransformerLayer
33from labml_nn.transformers.utils import subsequent_mask

#

带有 F TA 激活功能的 FF N 模块

36class FeedForwardFTA(nn.Module):

#

d_model 是令牌嵌入中的要素数量
d_ff 是 FFN 隐藏层中的要素数量
activation 是 FTA 激活模块
dropout 是隐藏层的丢失概率

41    def __init__(self, d_model: int, d_ff: int,
42                 activation: FTA,
43                 dropout: float = 0.1):

#

50        super().__init__()

#

第一层按权重 $W_{1}$ 和偏差进行参数化 $b_{1}$

52        self.layer1 = nn.Linear(d_model, d_ff)

#

第二层按权重 $W_{1}$ 和偏差进行参数化 $b_{1}$

54        self.layer2 = nn.Linear(d_ff * activation.expansion_factor, d_model)

#

隐藏图层丢失

56        self.dropout = nn.Dropout(dropout)

#

激活功能 $f$

58        self.activation = activation

#

60    def forward(self, x: torch.Tensor):

#

$f (x W_{1} + b_{1})$

62        x = self.activation(self.layer1(x))

#

申请退学

64        x = self.dropout(x)

#

66        return self.layer2(x)

#

自回归模型

这是一个自回归变压器模型，它使用前馈网络和（模糊平铺激活）（index.html）。

69class AutoregressiveTransformer(Module):

#

n_tokens 是词汇表中代币的数量
d_model 是嵌入的大小
n_layers 是变压器层的数量
layer 是层。我们在变压器上使用这个n_layers 副本。

77    def __init__(self, n_tokens: int, d_model: int, n_layers: int, layer: TransformerLayer):

#

84        super().__init__()

#

带n_layers 层的变压器

86        self.transformer_layers = nn.ModuleList([copy.deepcopy(layer) for _ in range(n_layers)])

#

令牌嵌入层

89        self.emb = nn.Embedding(n_tokens, d_model)

#

读出层

91        self.readout = nn.Linear(d_model, n_tokens)

#

掩码将在第一次调用时初始化

94        self.mask = None

#

x 是形状的输入标记[seq_len, batch_size]

96    def forward(self, x: torch.Tensor):

#

创建自动回归遮罩

101        if self.mask is None or self.mask.size(0) != len(x):

#

后续的掩码，将掩盖令牌以免看到未来的代币

103            self.mask = subsequent_mask(len(x)).to(x.device)

#

获取令牌嵌入

106        x = self.emb(x)

#

变压器编码

108        for layer in self.transformer_layers:
109            x = layer(x=x, mask=self.mask)

#

获取日志

111        x = self.readout(x)

#

返回结果

114        return x, None

#

配置

这继承自 NLPAutoRegressionConfigs

117class Configs(NLPAutoRegressionConfigs):

#

型号

126    model: AutoregressiveTransformer

#

层数

129    n_layers: int = 4

#

$α$ 对 $β$ 于 DeepNorm

132    deep_norm_alpha: float
133    deep_norm_beta: float

#

关注的头部数量

136    n_heads: int = 4

#

嵌入大小

138    d_model: int = 256

#

每个注意头的大小

140    d_k: int = 16

#

前馈图层大小

142    d_ff: int = 256

#

自贸区

145    fta_lower_limit: float = -1.
146    fta_upper_limit: float = +1.
147    fta_delta: float = 0.2
148    fta_eta: float = 0.05

#

初始化模型

151@option(Configs.model)
152def _model(c: Configs):

#

创建 FTA 激活模块

158    fta = FTA(c.fta_lower_limit, c.fta_upper_limit, c.fta_delta, c.fta_eta)

#

创建变压器。我们重复使用TransformerLayer 和MultiHeadAttention 实现。

162    m = AutoregressiveTransformer(c.n_tokens, c.d_model, c.n_layers,
163                                  TransformerLayer(d_model=c.d_model,
164                                                   feed_forward=FeedForwardFTA(d_model=c.d_model,
165                                                                               d_ff=c.d_ff,
166                                                                               activation=fta,
167                                                                               dropout=0.1),
168                                                   self_attn=MultiHeadAttention(c.n_heads, c.d_model,
169                                                                                dropout_prob=0.0),
170                                                   dropout_prob=0.0))

#

移到设备

173    return m.to(c.device)

#

创建并运行实验

176def main():

#

创建实验

181    experiment.create(name="fta", writers={'screen',  'comet', 'labml'})

#

创建配置

183    conf = Configs()

#

覆盖配置

185    experiment.configs(conf, {

#

使用角色等级分词器

187        'tokenizer': 'character',

#

提示分隔符为空

189        'prompt_separator': '',

#

开始采样提示

191        'prompt': 'It is ',

#

使用小莎士比亚数据集

193        'text': 'tiny_shakespeare',

#

使用上下文大小为 $256$

196        'seq_len': 256,

#

训练 32 个时代

198        'epochs': 32,

#

批量大小 $16$

200        'batch_size': 16,

#

在训练和验证之间切换每个纪元的 $10$ 次数

202        'inner_iterations': 10,

#

没有预热的 Adam 优化器

205        'optimizer.optimizer': 'Adam',
206        'optimizer.learning_rate': 3e-4,
207    })

#

设置用于保存和加载的模型

210    experiment.add_pytorch_models({'model': conf.model})

#

开始实验

213    with experiment.start():

#

跑步训练

215        conf.run()

#

219if __name__ == '__main__':
220    main()