#

深度规范实验

13import copy
14
15import torch
16import torch.nn as nn
17
18from labml import experiment
19from labml.configs import option
20from labml_helpers.module import Module
21from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
22from labml_nn.normalization.deep_norm import DeepNormTransformerLayer
23from labml_nn.transformers import MultiHeadAttention
24from labml_nn.transformers.feed_forward import FeedForward

#

自回归模型

这是一个使用 DeepNorm 的自回归变压器模型。

27class AutoregressiveTransformer(Module):

#

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

34    def __init__(self, n_tokens: int, d_model: int, n_layers: int, layer: DeepNormTransformerLayer):

#

41        super().__init__()

#

带n_layers 层的变压器

43        self.transformer = nn.Sequential(*[copy.deepcopy(layer) for _ in range(n_layers)])

#

令牌嵌入层

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

#

读出层

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

#

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

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

#

获取令牌嵌入

55        x = self.emb(x)

#

变压器编码

57        x = self.transformer(x)

#

获取日志

59        x = self.readout(x)

#

返回结果

62        return x, None

#

配置

这继承自 NLPAutoRegressionConfigs

65class Configs(NLPAutoRegressionConfigs):

#

型号

74    model: AutoregressiveTransformer

#

层数

77    n_layers: int = 32

#

$α$ 对 $β$ 于 DeepNorm

80    deep_norm_alpha: float
81    deep_norm_beta: float

#

关注的头部数量

84    n_heads: int = 4

#

嵌入大小

86    d_model: int = 64

#

每个注意头的大小

88    d_k: int = 16

#

计算 $α$

$α = (2 M)^{\frac{1}{4}}$

91@option(Configs.deep_norm_alpha)
92def _deep_norm_alpha(c: Configs):

#

98    return (2. * c.n_layers) ** (1. / 4.)

#

计算 $β$

$β = (8 M)^{- \frac{1}{4}}$

101@option(Configs.deep_norm_beta)
102def _deep_norm_beta(c: Configs):

#

108    return (8. * c.n_layers) ** -(1. / 4.)

#

初始化模型

111@option(Configs.model)
112def _model(c: Configs):

#

116    m = AutoregressiveTransformer(c.n_tokens, c.d_model, c.n_layers,
117                                  DeepNormTransformerLayer(d_model=c.d_model,
118                                                           deep_norm_alpha=c.deep_norm_alpha,
119                                                           deep_norm_beta=c.deep_norm_beta,
120                                                           feed_forward=FeedForward(d_model=c.d_model,
121                                                                                    d_ff=c.d_model * 4),
122                                                           self_attn=MultiHeadAttention(c.n_heads, c.d_model,
123                                                                                        dropout_prob=0.0)))
124
125    return m.to(c.device)

#

创建并运行实验

128def main():

#

创建实验

133    experiment.create(name="deep_norm", writers={'screen', 'web_api'})

#

创建配置

135    conf = Configs()

#

覆盖配置

137    experiment.configs(conf, {

#

使用角色等级分词器

139        'tokenizer': 'character',

#

提示分隔符为空

141        'prompt_separator': '',

#

开始采样提示

143        'prompt': 'It is ',

#

使用小莎士比亚数据集

145        'text': 'tiny_shakespeare',

#

使用上下文大小为 $256$

148        'seq_len': 256,

#

训练 32 个时代

150        'epochs': 32,

#

批量大小 $16$

152        'batch_size': 16,

#

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

154        'inner_iterations': 10,

#

层数

157        'n_layers': 50,

#

没有预热的 Adam 优化器

161        'optimizer.optimizer': 'Adam',
162        'optimizer.learning_rate': 1.25e-4,
163    })

#

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

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

#

开始实验

169    with experiment.start():

#

跑步训练

171        conf.run()

#

175if __name__ == '__main__':
176    main()