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

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

35class FeedForwardFTA(nn.Module):

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

49        super().__init__()

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

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

55        self.dropout = nn.Dropout(dropout)

57        self.activation = activation

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

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

63        x = self.dropout(x)

65        return self.layer2(x)

68class AutoregressiveTransformer(Module):

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

83        super().__init__()

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

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

116class Configs(NLPAutoRegressionConfigs):

125    model: AutoregressiveTransformer

128    n_layers: int = 4

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

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

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

195        'seq_len': 256,

197        'epochs': 32,

199        'batch_size': 16,

201        'inner_iterations': 10,

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