# 分层转换器是更有效的语言模型实验

14import math
15from typing import List
16
17import torch
18from torch import nn
19
20from labml import experiment
21from labml.configs import option
22from labml_helpers.module import Module
23from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
24from labml_nn.transformers.hour_glass import HourGlass
25from labml_nn.transformers.positional_encoding import PositionalEncoding

## 自回归语言模型

28class AutoregressiveTransformer(Module):
• n_tokens 词汇量是多少
• d_model 是令牌嵌入的大小
• dropout 是辍学概率
• hour_glass沙漏型号吗
33    def __init__(self, n_tokens: int, d_model: int, dropout: float, hour_glass: HourGlass):
40        super().__init__()

42        self.embedding = nn.Embedding(n_tokens, d_model)
48        self.pos_embedding = PositionalEncoding(d_model, dropout)
50        self.hour_glass = hour_glass

52        self.norm = nn.LayerNorm([d_model])

54        self.d_model = d_model

56        self.output = nn.Linear(d_model, n_tokens)
• x 是令牌索引为形状的张量[seq_len, batch_size]
58    def __call__(self, x: torch.Tensor):

63        x = self.embedding(x)

66        if self.pos_embedding is not None:
67            x = self.pos_embedding(x * math.sqrt(self.d_model))

70        x = self.hour_glass(x)

73        output = self.output(self.norm(x))

76        return output, None

## 配置

79class Configs(NLPAutoRegressionConfigs):

87    model: AutoregressiveTransformer

89    n_heads: int = 8

91    dropout: float = 0.1

93    d_ff: int = 512

95    d_model: int = 256

97    shortening_factors: List[int] = [8, 4]

100@option(Configs.model)
101def _model(c: Configs):

107    hour_glass = HourGlass(c.n_heads, c.d_model, c.dropout, c.d_ff, c.shortening_factors)

109    m = AutoregressiveTransformer(c.n_tokens, c.d_model, c.dropout, hour_glass).to(c.device)
112    return m
115def main():

117    experiment.create(name="hour_glass")

119    conf = Configs()

121    experiment.configs(conf, {

123        'tokenizer': 'character',

125        'prompt_separator': '',

127        'prompt': 'It is ',

129        'text': 'tiny_shakespeare',

132        'seq_len': 256,

134        'epochs': 128,

136        'batch_size': 32,

139        'inner_iterations': 10,
142        'optimizer.optimizer': 'Noam',
143        'optimizer.learning_rate': 1.,
145    })

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

151    with experiment.start():

153        conf.run()
157if __name__ == '__main__':
158    main()