#

模糊拼贴激活实验

在这里，我们训练一台在前馈网络中使用模糊切片激活的变压器。我们将其用作语言模型，并在小莎士比亚数据集上对其进行训练以进行演示。

但是，对于 FTA 来说，这可能不是理想的任务，我们认为 FTA 更适合对具有连续变量的数据进行建模。

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

#

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

35class FeedForwardFTA(nn.Module):

#

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

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

#

49        super().__init__()

#

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

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

#

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

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

#

隐藏图层丢失

55        self.dropout = nn.Dropout(dropout)

#

激活功能 $f$

57        self.activation = activation

#

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

#

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

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

#

申请退学

63        x = self.dropout(x)

#

65        return self.layer2(x)

#

自回归模型

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

68class AutoregressiveTransformer(Module):

#

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

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

#

83        super().__init__()

#

带n_layers 层的变压器

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

#

令牌嵌入层

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

#

读出层

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

#

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

93        self.mask = None

#

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

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

#

创建自动回归遮罩

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

#

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

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

#

获取令牌嵌入

105        x = self.emb(x)

#

变压器编码

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

#

获取日志

110        x = self.readout(x)

#

返回结果

113        return x, None

#

配置

这继承自 NLPAutoRegressionConfigs

116class Configs(NLPAutoRegressionConfigs):

#

型号

125    model: AutoregressiveTransformer

#

层数

128    n_layers: int = 4

#

$α$ 对 $β$ 于 DeepNorm

131    deep_norm_alpha: float
132    deep_norm_beta: float

#

关注的头部数量

135    n_heads: int = 4

#

嵌入大小

137    d_model: int = 256

#

每个注意头的大小

139    d_k: int = 16

#

前馈图层大小

141    d_ff: int = 256

#

自贸区

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

#

初始化模型

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

#

创建 FTA 激活模块

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

#

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

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

#

移到设备

172    return m.to(c.device)

#

创建并运行实验

175def main():

#

创建实验

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

#

创建配置

182    conf = Configs()

#

覆盖配置

184    experiment.configs(conf, {

#

使用角色等级分词器

186        'tokenizer': 'character',

#

提示分隔符为空

188        'prompt_separator': '',

#

开始采样提示

190        'prompt': 'It is ',

#

使用小莎士比亚数据集

192        'text': 'tiny_shakespeare',

#

使用上下文大小为 $256$

195        'seq_len': 256,

#

训练 32 个时代

197        'epochs': 32,

#

批量大小 $16$

199        'batch_size': 16,

#

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

201        'inner_iterations': 10,

#

没有预热的 Adam 优化器

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

#

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

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

#

开始实验

212    with experiment.start():

#

跑步训练

214        conf.run()

#

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