ගැඹුරු සම්මතය අත්හදා බැලීම

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

27class AutoregressiveTransformer(Module):

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

41        super().__init__()

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)

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

65class Configs(NLPAutoRegressionConfigs):

74    model: AutoregressiveTransformer

77    n_layers: int = 32

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

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

98    return (2. * c.n_layers) ** (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',

148        'seq_len': 256,

150        'epochs': 32,

152        'batch_size': 16,

154        'inner_iterations': 10,

157        'n_layers': 50,

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()