සම්පීඩ්යතාට්රාන්ස්ෆෝමර් අත්හදා බැලීම

11from typing import List, Tuple, NamedTuple
12
13import torch
14import torch.nn as nn
15
16from labml import experiment, tracker, monit, logger
17from labml.configs import option
18from labml.logger import Text
19from labml_helpers.metrics.simple_state import SimpleStateModule
20from labml_helpers.module import Module
21from labml_helpers.train_valid import BatchIndex, hook_model_outputs
22from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
23from labml_nn.transformers.compressive import CompressiveTransformer, AttentionReconstructionLoss, \
24    CompressiveTransformerLayer, Conv1dCompression

27class CompressedMemory(NamedTuple):
28    mem: List[torch.Tensor]
29    c_mem: List[torch.Tensor]

32class AutoregressiveModel(Module):

37    def __init__(self, n_vocab: int, d_model: int, transformer: CompressiveTransformer):
38        super().__init__()

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

42        self.transformer = transformer

44        self.generator = nn.Linear(d_model, n_vocab)

46        self.mask_x = None
47        self.mask_mem = None

49    def forward(self, x: torch.Tensor, mem: CompressedMemory):

51        if mem is not None:
52            mem, c_mem = mem.mem, mem.c_mem
53        else:
54            mem = []
55            c_mem = []

58        m_len = len(mem[0]) if mem else 0
59        if c_mem:
60            m_len += len(c_mem[0])

63        if self.mask_x is None or self.mask_x.shape[0] < len(x):
64            from labml_nn.transformers.utils import subsequent_mask
65            self.mask_x = subsequent_mask(len(x)).to(x.device)

67        if self.mask_mem is None or self.mask_mem.shape[1] < m_len or self.mask_mem.shape[0] < len(x):
68            self.mask_mem = self.mask_x.new_ones(len(x), m_len, 1)

71        if m_len:
72            mask = torch.cat((self.mask_mem[:len(x), :m_len], self.mask_x[:len(x), :len(x)]), dim=1)

74        else:
75            mask = self.mask_x[:len(x), :len(x)]

78        x = self.src_embed(x)

80        res, mem = self.transformer(x, mem, c_mem, mask)

82        res = self.generator(res)

84        return res, mem

87class Configs(NLPAutoRegressionConfigs):

94    model: AutoregressiveModel

97    d_model: int = 128

99    heads: int = 4

101    dropout: float = 0.0

103    d_ff: int = 256

105    n_layers: int = 6

107    mem_len: int = 8

109    memory = SimpleStateModule()

111    attention_reconstruction_loss: AttentionReconstructionLoss

113    compression_rate: int = 4

115    c_mem_len: int = 128

117    def init(self):

119        tracker.set_scalar("accuracy.*", True)
120        tracker.set_scalar("loss.*", True)

122        tracker.set_scalar("ar_loss.*", False)

124        hook_model_outputs(self.mode, self.model, 'model')

126        self.state_modules = [self.accuracy, self.memory]

128    @torch.no_grad()
129    def merge_compress_memory(self, mem: CompressedMemory, new_mem: List[torch.Tensor]) \
130            -> Tuple[CompressedMemory, List[torch.Tensor]]:

136        if self.mem_len == 0 and self.c_mem_len == 0:
137            return CompressedMemory([], []), []

140        if mem is not None:
141            mem, c_mem = mem.mem, mem.c_mem
142        else:
143            mem, c_mem = [], []

146        if mem:
147            mem = [torch.cat((m, x), dim=0) for m, x in zip(mem, new_mem)]
148        else:
149            mem = new_mem

152        if len(mem[0]) > self.mem_len:

156            n_c_mem = (len(mem[0]) - self.mem_len + self.compression_rate - 1) // self.compression_rate

158            n_old = n_c_mem * self.compression_rate

160            mem_to_compress = []

162            uncompressed_mem = []

164            for m in mem:

166                cm, m = torch.split(m, [n_old, len(m) - n_old])

168                mem_to_compress.append(cm)

170                uncompressed_mem.append(m)

172            mem = uncompressed_mem

175            new_c_mem = []
176            for i, layer in enumerate(self.model.transformer.layers):
177                new_c_mem.append(layer.compress(mem_to_compress[i]))

180            if c_mem:
181                c_mem = [torch.cat((m, nm), dim=0) for m, nm in zip(c_mem, new_c_mem)]

183            else:
184                c_mem = new_c_mem

187            if len(c_mem[0]) > self.c_mem_len:
188                c_mem = [m[-self.c_mem_len:] for m in c_mem]

190        else:
191            mem_to_compress = []

195        return CompressedMemory(mem, c_mem), mem_to_compress

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

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

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

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

212            mem = self.memory.get()

214            output, new_mem = self.model(data, mem)

216            mem, mem_to_compress = self.merge_compress_memory(mem, new_mem)

218            self.memory.set(mem)

221        loss = self.loss_func(output, target)
222        tracker.add("loss.", loss)

225        if mem_to_compress:

227            ar_loss = self.attention_reconstruction_loss(new_mem, mem_to_compress)

229            tracker.add("ar_loss.", ar_loss)

231            loss = loss + ar_loss

234        self.accuracy(output, target)
235        self.accuracy.track()

238        if self.mode.is_train:

240            loss.backward()

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

244            self.optimizer.step()

246            if batch_idx.is_last:
247                tracker.add('model', self.model)

249            self.optimizer.zero_grad()

252        tracker.save()

254    def sample(self):

260        prompt = self.prompt

262        log = [(prompt, Text.subtle)]

264        mem = CompressedMemory([], [])

266        for i in monit.iterate('Sample', 25):

268            data = self.text.text_to_i(prompt).unsqueeze(-1)

270            data = data.to(self.device)

272            output, new_mem = self.model(data, mem)

274            output = output.argmax(dim=-1).squeeze(1)

276            prompt += self.prompt_separator + self.text.itos[output[-1]]

278            prompt = prompt[-1:]

280            log += [(self.prompt_separator + self.text.itos[output[-1]], Text.value)]

282            mem, _ = self.merge_compress_memory(mem, new_mem)

285        logger.log(log)

288@option(Configs.model)
289def autoregressive_model(c: Configs):

293    from labml_nn.transformers.xl import RelativeMultiHeadAttention
294    from labml_nn.transformers.feed_forward import FeedForward
295    m = AutoregressiveModel(c.n_tokens, c.d_model, CompressiveTransformer(
296        CompressiveTransformerLayer(d_model=c.d_model,
297                                    self_attn=RelativeMultiHeadAttention(c.heads, c.d_model, c.dropout),
298                                    feed_forward=FeedForward(c.d_model, c.d_ff, c.dropout),
299                                    dropout_prob=c.dropout,
300                                    compress=Conv1dCompression(c.compression_rate, c.d_model)), c.n_layers))
301    return m.to(c.device)

304@option(Configs.attention_reconstruction_loss)
305def attention_reconstruction_loss(c: Configs):

309    return AttentionReconstructionLoss(c.model.transformer.layers)

312def main():

317    experiment.create(name="compressive_transformer", comment='')

319    conf = Configs()

321    experiment.configs(conf,

323                       {'tokenizer': 'character',
324                        'text': 'tiny_shakespeare',
325                        'optimizer.learning_rate': 2.5e-4,
326                        'optimizer.optimizer': 'AdamW',
327                        'prompt': 'It is',
328                        'prompt_separator': '',
329
330                        'train_loader': 'sequential_train_loader',
331                        'valid_loader': 'sequential_valid_loader',
332
333                        'seq_len': 8,
334                        'mem_len': 8,
335                        'epochs': 128,
336                        'batch_size': 32,
337                        'inner_iterations': 25,
338                        'compression_rate': 2,
339                        })

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

345    with experiment.start():

347        conf.run()

351if __name__ == '__main__':
352    main()