#

复古训练

这是 RETRO 的训练代码。

14import torch
15from torch import nn
16from torch.utils.data import DataLoader, RandomSampler
17
18from labml import monit, lab, tracker, experiment, logger
19from labml.logger import Text
20from labml_helpers.datasets.text import TextFileDataset
21from labml_nn.optimizers.noam import Noam
22from labml_nn.transformers.retro import model as retro
23from labml_nn.transformers.retro.dataset import Dataset, RetroIndex
24from labml_nn.transformers.retro.model import RetroModel, NearestNeighborEncoder

#

采样器

这个类贪婪地从模型中抽样。

27class Sampler:

#

device 是该型号的设备
model 是复古模式
tds 是文本数据集（用于获取相邻数据块）
chunk_len 是区块的长度

34    def __init__(self, device: torch.device, model: retro.RetroModel, tds: TextFileDataset, chunk_len: int):

#

41        self.chunk_len = chunk_len
42        self.tds = tds
43        self.model = model
44        self.device = device

#

复古指数

47        self.index = RetroIndex()

#

检索给定区块的最近邻域

49    def retrieve_nearest_neighbours(self, chunk: str):

#

检索最近邻点的偏移量

55        neighbor_offsets = self.index([chunk], None)

#

获取邻居（邻居长度等于chunk_len * 2 ）

58        text = self.tds.train
59        neighbors = [text[j: j + self.chunk_len * 2] for j in neighbor_offsets[0]]

#

62        return neighbors

#

来自给定提示的示例文本

64    def sample(self, prompt: str, sample_len: int):

#

将最近邻存储为字符串

70        neighbors_str = []

#

样本文本

73        sampled = ''

#

sample_len 代币样本

76        for i in range(sample_len):

#

如果采样块比我们已经检索到的要多，我们就需要检索邻居

79            while len(neighbors_str) < len(prompt) // self.chunk_len:

#

获取我们尚未检索到邻居的最后一个区块

81                off = len(neighbors_str) * self.chunk_len
82                chunk = prompt[off: off + self.chunk_len]

#

检索最近的邻居

84                neighbors_str.append(self.retrieve_nearest_neighbours(chunk))

#

对输入进行标记化

87            src = self.tds.text_to_i(prompt)

#

标记检索到的邻居

89            neighbors = torch.stack([torch.stack([self.tds.text_to_i(n) for n in chunk]) for chunk in neighbors_str])

#

将它们移到与模型相同的设备上

92            src = src.to(self.device)
93            neighbors = neighbors.to(self.device)

#

获取模型输出

96            res = self.model(src[None, :], neighbors[None, :, :, :])

#

贪婪地抽取最后一个代币

99            token = res[0, -1, :].argmax(dim=-1)

#

将采样标记文本添加到提示和示例文本中

102            prompt += self.tds.itos[token.item()]
103            sampled += self.tds.itos[token.item()]

#

106        return sampled

#

复古教练

109class Trainer:

#

device 是该型号的设备
model 是复古模式
dataloader 是具有预检索邻域的数据集的数据加载器
optimizer 是优化器

114    def __init__(self, device: torch.device, model: retro.RetroModel,
115                 dataloader: DataLoader, optimizer: torch.optim.Optimizer):

#

122        self.optimizer = optimizer
123        self.device = device
124        self.dataloader = dataloader
125        self.model = model
126        self.loss_func = nn.CrossEntropyLoss()

#

为一个时代训练模型

128    def __call__(self):

#

遍历训练数据

134        for i, (src, tgt, neighbors) in monit.enum('Train', self.dataloader):

#

将数据移动到设备

136            src, tgt, neighbors = src.to(self.device), tgt.to(self.device), neighbors.to(self.device)

#

向前传球

139            res = self.model(src, neighbors)

#

计算损失

141            loss = self.loss_func(res.view(-1, res.shape[-1]), tgt.view(-1))

#

清除渐变

144            self.optimizer.zero_grad()

#

向后传球

146            loss.backward()

#

优化模型

148            self.optimizer.step()

#

保存训练统计数据并增加全局步数计数器

151            tracker.save({'loss.train': loss})
152            tracker.add_global_step(len(src))

#

创建和训练一个小模型

155def train():

#

创建实验

161    experiment.create(name='retro_small')

#

GPU 设备

164    device = torch.device('cuda:0')

#

加载小莎士比亚数据集

167    tds = TextFileDataset(
168        lab.get_data_path() / 'tiny_shakespeare.txt',
169        list,
170        url='https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt')

#

加载复古数据集

173    train_dataset = Dataset(lab.get_data_path() / 'retro_train_dataset.json', tds)

#

创建数据加载器

176    train_dl = DataLoader(train_dataset,
177                          batch_size=4,
178                          sampler=RandomSampler(train_dataset, replacement=True))

#

超参数

181    chunk_len = 16
182    d_model = 128
183    d_ff = 512
184    n_heads = 16
185    d_k = 16

#

创建最近邻编码器

188    nearest_neighbor_encoder = NearestNeighborEncoder(chunk_len, 6, {3}, d_model, n_heads, d_k, d_ff)

#

创建模型

190    model = RetroModel(tds.n_tokens, d_model, 6,
191                       {3, 5},
192                       chunk_len, n_heads, d_k, d_ff,
193                       encoder=nearest_neighbor_encoder)

#

将模型移到设备上

195    model = model.to(device)

#

创建优化器

197    optimizer = Noam(model.parameters(), lr=1., d_model=d_model, warmup=2_000)

#

创建Trainer

199    trainer = Trainer(device, model, train_dl, optimizer)

#

创建Sampler

201    sampler = Sampler(device, model, tds, chunk_len)

#

203    prompt = '''Second Citizen:\nOne word, good citizens.\n\nFirst Citizen:'''

#

设置用于保存和加载的模型

206    experiment.add_pytorch_models(model=model)

#

开始实验

209    with experiment.start():

#

为32 时代而训练

211        for epoch in monit.loop(32):

#

火车

213            trainer()

#

打印一条新行

215            tracker.new_line()

#

样本来自prompt

217            logger.log([(prompt.replace('\n', '\\n\n'), Text.subtle),
218                        (sampler.sample(prompt, 128).replace('\n', '\\n\n'), Text.none)])

#

保存模型

220            experiment.save_checkpoint()

#

224if __name__ == '__main__':
225    train()