නොපැහැදිලි ටයිල් සක්රීය කිරීමේ අත්හදා බැලීම

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

36class FeedForwardFTA(nn.Module):

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

50        super().__init__()

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

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

56        self.dropout = nn.Dropout(dropout)

58        self.activation = activation

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

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

64        x = self.dropout(x)

66        return self.layer2(x)

69class AutoregressiveTransformer(Module):

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

84        super().__init__()

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

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

117class Configs(NLPAutoRegressionConfigs):

126    model: AutoregressiveTransformer

129    n_layers: int = 4

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

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

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',

196        'seq_len': 256,

198        'epochs': 32,

200        'batch_size': 16,

202        'inner_iterations': 10,

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