ස්විච්ට්රාන්ස්ෆෝමර් අත්හදා බැලීම

15import torch
16import torch.nn as nn
17
18from labml import experiment, tracker
19from labml.configs import option
20from labml_helpers.module import Module
21from labml_helpers.train_valid import BatchIndex
22from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs

25class AutoregressiveModel(Module):

30    def __init__(self, n_vocab: int, d_model: int, transformer: Module):
31        super().__init__()

33        self.src_embed = nn.Embedding(n_vocab, d_model)

35        self.transformer = transformer

37        self.generator = nn.Linear(d_model, n_vocab)
38        self.mask = None

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

42        if self.mask is None or self.mask.size(0) != len(x):
43            from labml_nn.transformers.utils import subsequent_mask
44            self.mask = subsequent_mask(len(x)).to(x.device)

46        x = self.src_embed(x)

48        res, counts, route_prob, n_dropped, route_prob_max = self.transformer(x, self.mask)

50        res = self.generator(res)

52        return res, counts, route_prob, n_dropped, route_prob_max

55class Configs(NLPAutoRegressionConfigs):

64    model: AutoregressiveModel
65    transformer: Module

68    d_model: int = 128

70    heads: int = 4

72    dropout: float = 0.0

74    d_ff: int = 256

76    n_layers: int = 6

78    n_experts: int = 4

80    load_balancing_loss_ceof = 0.01

82    is_scale_prob: bool = True

84    drop_tokens: bool = False

86    capacity_factor: float = 1.0

88    def init(self):
89        super().init()

91        tracker.set_scalar("lb_loss.*", False)
92        tracker.set_scalar("route.*", False)
93        tracker.set_scalar("dropped.*", False)

95    def step(self, batch: any, batch_idx: BatchIndex):

101        data, target = batch[0].to(self.device), batch[1].to(self.device)

104        if self.mode.is_train:
105            tracker.add_global_step(data.shape[0] * data.shape[1])

108        with self.mode.update(is_log_activations=batch_idx.is_last):

110            output, counts, route_prob, n_dropped, route_prob_max = self.model(data)

113        cross_entropy_loss = self.loss_func(output, target)

115        total = counts.sum(dim=-1, keepdims=True)

119        route_frac = counts / total

122        route_prob = route_prob / total

127        load_balancing_loss = self.n_experts * (route_frac * route_prob).sum()

130        tracker.add('dropped.', total.new_tensor(n_dropped) / total)
131        tracker.add('route.min.', route_frac.min())
132        tracker.add('route.max.', route_frac.max())
133        tracker.add('route.std.', route_frac.std())
134        tracker.add('route.max_prob.', route_prob_max)
135        tracker.add("loss.", cross_entropy_loss)
136        tracker.add("lb_loss.", load_balancing_loss)

141        loss = cross_entropy_loss + self.load_balancing_loss_ceof * load_balancing_loss

144        self.accuracy(output, target)
145        self.accuracy.track()

148        if self.mode.is_train:

150            loss.backward()

152            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.grad_norm_clip)

154            self.optimizer.step()

156            if batch_idx.is_last:
157                tracker.add('model', self.model)

159            self.optimizer.zero_grad()

162        tracker.save()

165@option(Configs.model)
166def autoregressive_model(c: Configs):

170    m = AutoregressiveModel(c.n_tokens, c.d_model, c.transformer)
171    return m.to(c.device)

174@option(Configs.transformer)
175def switch_transformer(c: Configs):

179    from labml_nn.transformers.switch import SwitchTransformer, SwitchTransformerLayer, SwitchFeedForward
180    from labml_nn.transformers import MultiHeadAttention
181    from labml_nn.transformers.feed_forward import FeedForward
182
183    return SwitchTransformer(
184        SwitchTransformerLayer(d_model=c.d_model,
185                               attn=MultiHeadAttention(c.heads, c.d_model, c.dropout),
186                               feed_forward=SwitchFeedForward(capacity_factor=c.capacity_factor,
187                                                              drop_tokens=c.drop_tokens,
188                                                              is_scale_prob=c.is_scale_prob,
189                                                              n_experts=c.n_experts,
190                                                              expert=FeedForward(c.d_model, c.d_ff, c.dropout),
191                                                              d_model=c.d_model),
192                               dropout_prob=c.dropout),
193        c.n_layers)

196def main():

201    experiment.create(name="switch_transformer", comment='')

203    conf = Configs()

205    experiment.configs(conf,

207                       {'tokenizer': 'character',
208                        'text': 'tiny_shakespeare',
209                        'optimizer.learning_rate': 1.,
210                        'optimizer.optimizer': 'Noam',
211                        'prompt': 'It is',
212                        'prompt_separator': '',
213
214                        'transformer': 'switch_transformer',
215                        'n_experts': 4,
216
217                        'drop_tokens': True,
218                        'capacity_factor': 1.2,
219
220                        'train_loader': 'shuffled_train_loader',
221                        'valid_loader': 'shuffled_valid_loader',
222
223                        'seq_len': 64,
224                        'epochs': 128,
225                        'batch_size': 32,
226                        'inner_iterations': 25,
227                        })

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

233    with experiment.start():

235        conf.run()

239if __name__ == '__main__':
240    main()