从头开始构建和训练 Transformer（下）

导读

上一篇推文从头开始构建和训练 Transformer（上）https://blog.csdn.net/weixin_46287760/article/details/136048418介绍了构建和训练Transformer的过程和构建每个组件的代码示例。本文将使用数据对该架构进行代码演示，验证其模型性能。

本期『数据+代码』已上传百度网盘。

有需要的朋友关注公众号【小Z的科研日常】，回复关键词[Transformer]获取。

01、加载数据集

对于此任务，我们将使用🤗Hugging Face 上提供的OpusBooks 数据集。该数据集由两个特征组成，id和translation。该translation功能包含不同语言的句子对，例如西班牙语和葡萄牙语、英语和法语等。

我首先尝试将句子从英语翻译成葡萄牙语，但是这对句子只有 1.4k个示例，因此在该模型的当前配置中结果并不令人满意。然后，我尝试使用英语-法语对，因为它的示例数量较多（127k），但使用当前配置进行训练需要很长时间。

我们首先定义get_all_sentences函数来迭代数据集并根据定义的语言对提取句子。

# 迭代数据集，提取原句及其译文
def get_all_sentences(ds, lang):
    for pair in ds:
        yield pair['translation'][lang]

该get_ds函数定义为加载和准备数据集以进行训练和验证。在此函数中，我们构建或加载分词器、拆分数据集并创建 DataLoader，以便模型可以成功地批量迭代数据集。这些函数的结果是源语言和目标语言的标记器以及 DataLoader 对象。

def get_ds(config):
    
      # 语言对将在我们稍后创建的 "配置 "字典中定义。
    ds_raw = load_dataset('opus_books', f'{config["lang_src"]}-{config["lang_tgt"]}', split = 'train') 
    
    # 为源语言和目标语言构建或加载标记符
    tokenizer_src = build_tokenizer(config, ds_raw, config['lang_src'])
    tokenizer_tgt = build_tokenizer(config, ds_raw, config['lang_tgt'])
    
    # 分割数据集进行训练和验证
    train_ds_size = int(0.9 * len(ds_raw)) # 90% for training
    val_ds_size = len(ds_raw) - train_ds_size # 10% for validation
    train_ds_raw, val_ds_raw = random_split(ds_raw, [train_ds_size, val_ds_size]) # Randomly splitting the dataset
                                    
    # 使用双语数据集（BilingualDataset）类处理数据，我们将在下面定义该类
    train_ds = BilingualDataset(train_ds_raw, tokenizer_src, tokenizer_tgt, config['lang_src'], config['lang_tgt'], config['seq_len'])
    val_ds = BilingualDataset(val_ds_raw, tokenizer_src, tokenizer_tgt, config['lang_src'], config['lang_tgt'], config['seq_len'])
                                    
   # 对整个数据集进行迭代，并打印在源语言和目标语言句子中找到的最大长度
    max_len_src = 0
    max_len_tgt = 0
    for pair in ds_raw:
        src_ids = tokenizer_src.encode(pair['translation'][config['lang_src']]).ids
        tgt_ids = tokenizer_src.encode(pair['translation'][config['lang_tgt']]).ids
        max_len_src = max(max_len_src, len(src_ids))
        max_len_tgt = max(max_len_tgt, len(tgt_ids))
        
    print(f'Max length of source sentence: {max_len_src}')
    print(f'Max length of target sentence: {max_len_tgt}')
    
    # 为训练集和验证集创建数据加载器
    # 在训练和验证过程中，使用数据加载器分批迭代数据集
    train_dataloader = DataLoader(train_ds, batch_size = config['batch_size'], shuffle = True) # Batch size will be defined in the config dictionary
    val_dataloader = DataLoader(val_ds, batch_size = 1, shuffle = True)
    
    return train_dataloader, val_dataloader, tokenizer_src, tokenizer_tgt # Returning the DataLoader objects and tokenizers

我们定义casual_mask函数来为解码器的注意力机制创建掩码。此掩码可防止模型获得有关序列中未来元素的信息。

我们首先制作一个充满 1 的方形网格。我们用参数确定网格大小size。然后，我们将主对角线上方的所有数字更改为零。一侧的每个数字都变成零，而其余的仍然是1。然后该函数翻转所有这些值，将 1 变为 0，将 0 变为 1。这个过程对于预测序列中未来标记的模型至关重要。

02、验证循环

我们现在将为验证循环创建两个函数。验证循环对于评估模型从训练期间未见过的数据翻译句子的性能至关重要。

我们将定义两个函数。第一个函数，greedy_decode通过获取最可能的下一个标记为我们提供模型的输出。第二个函数run_validation负责运行验证过程，在该过程中我们解码模型的输出并将其与目标句子的参考文本进行比较。

class BilingualDataset(Dataset):
    def __init__(self, ds, tokenizer_src, tokenizer_tgt, src_lang, tgt_lang, seq_len) -> None:
        super().__init__()
        
        self.seq_len = seq_len
        self.ds = ds
        self.tokenizer_src = tokenizer_src
        self.tokenizer_tgt = tokenizer_tgt
        self.src_lang = src_lang
        self.tgt_lang = tgt_lang
        
        
        self.sos_token = torch.tensor([tokenizer_tgt.token_to_id("[SOS]")], dtype=torch.int64)
        self.eos_token = torch.tensor([tokenizer_tgt.token_to_id("[EOS]")], dtype=torch.int64)
        self.pad_token = torch.tensor([tokenizer_tgt.token_to_id("[PAD]")], dtype=torch.int64)

        
    
    def __len__(self):
        return len(self.ds)
    
    
    def __getitem__(self, index: Any) -> Any:
        src_target_pair = self.ds[index]
        src_text = src_target_pair['translation'][self.src_lang]
        tgt_text = src_target_pair['translation'][self.tgt_lang]
        
       
        enc_input_tokens = self.tokenizer_src.encode(src_text).ids
        dec_input_tokens = self.tokenizer_tgt.encode(tgt_text).ids
        
       
        enc_num_padding_tokens = self.seq_len - len(enc_input_tokens) - 2 # Subtracting the two '[EOS]' and '[SOS]' special tokens
       
        dec_num_padding_tokens = self.seq_len - len(dec_input_tokens) - 1 # Subtracting the '[SOS]' special token
        
        if enc_num_padding_tokens < 0 or dec_num_padding_tokens < 0:
            raise ValueError('Sentence is too long')
         
        encoder_input = torch.cat(
            [
            self.sos_token, # inserting the '[SOS]' token
            torch.tensor(enc_input_tokens, dtype = torch.int64), # Inserting the tokenized source text
            self.eos_token, # Inserting the '[EOS]' token
            torch.tensor([self.pad_token] * enc_num_padding_tokens, dtype = torch.int64) # Addind padding tokens
            ]
        )
        
        decoder_input = torch.cat(
            [
                self.sos_token, # inserting the '[SOS]' token 
                torch.tensor(dec_input_tokens, dtype = torch.int64), # Inserting the tokenized target text
                torch.tensor([self.pad_token] * dec_num_padding_tokens, dtype = torch.int64) # Addind padding tokens
            ]
        
        )

        label = torch.cat(
            [
                torch.tensor(dec_input_tokens, dtype = torch.int64), # Inserting the tokenized target text
                self.eos_token, # Inserting the '[EOS]' token 
                torch.tensor([self.pad_token] * dec_num_padding_tokens, dtype = torch.int64) # Adding padding tokens
                
            ]
        )

        assert encoder_input.size(0) == self.seq_len
        assert decoder_input.size(0) == self.seq_len
        assert label.size(0) == self.seq_len
        
        return {
            'encoder_input': encoder_input,
            'decoder_input': decoder_input, 
            'encoder_mask': (encoder_input != self.pad_token).unsqueeze(0).unsqueeze(0).int(),
            'decoder_mask': (decoder_input != self.pad_token).unsqueeze(0).unsqueeze(0).int() & casual_mask(decoder_input.size(0)), 
            'label': label,
            'src_text': src_text,
            'tgt_text': tgt_text
        }

03、训练循环

我们已准备好在 OpusBook 数据集上训练 Transformer 模型，以执行英语到意大利语翻译任务。我们首先通过调用我们之前定义的get_model函数来定义加载模型的函数。build_transformer该函数使用config字典来设置一些参数。

def get_model(config, vocab_src_len, vocab_tgt_len):
    
    model = build_transformer(vocab_src_len, vocab_tgt_len, config['seq_len'], config['seq_len'], config['d_model'])
    return model

下面我们将定义两个函数来配置我们的模型和训练过程。

在get_config函数中，我们定义了训练过程的关键参数。batch_size一次迭代中使用的训练示例的数量、num_epochs整个数据集通过 Transformer 向前和向后传递的次数、lr优化器的学习率等。我们最终还将定义来自 OpusBook 数据集的对，'lang_src': 'en'用于选择英语作为源语言以及'lang_tgt': 'it'选择意大利语作为目标语言。

该get_weights_file_path函数构建用于保存或加载任何特定时期的模型权重的文件路径。

def get_config():
    return{
        'batch_size': 8,
        'num_epochs': 20,
        'lr': 10**-4,
        'seq_len': 350,
        'd_model': 512, 
        'lang_src': 'en',
        'lang_tgt': 'it',
        'model_folder': 'weights',
        'model_basename': 'tmodel_',
        'preload': None,
        'tokenizer_file': 'tokenizer_{0}.json',
        'experiment_name': 'runs/tmodel'
    }

def get_weights_file_path(config, epoch: str):
    model_folder = config['model_folder'] 
    model_basename = config['model_basename'] 
    model_filename = f"{model_basename}{epoch}.pt" 
    return str(Path('.')/ model_folder/ model_filename)

我们最终定义了最后一个函数，train_model它将config参数作为输入。

在此函数中，我们将为训练设置一切。我们将模型及其必要组件加载到 GPU 上以加快训练速度，设置Adam优化器并配置CrossEntropyLoss函数来计算模型输出的翻译与数据集中的参考翻译之间的差异。

迭代训练批次、执行反向传播和计算梯度所需的每个循环都在此函数中。我们还将使用它来运行验证函数并保存模型的当前状态。

def train_model(config):
    # 设置设备在 GPU 上运行，以加快训练速度
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device {device}")
    
    # 创建模型目录以存储权重
    Path(config['model_folder']).mkdir(parents=True, exist_ok=True)
    
    # 使用 "get_ds "函数检索源语言和目标语言的数据加载器和标记器
    train_dataloader, val_dataloader, tokenizer_src, tokenizer_tgt = get_ds(config)
    
    # 使用 "get_model "函数在 GPU 上初始化模型
    model = get_model(config,tokenizer_src.get_vocab_size(), tokenizer_tgt.get_vocab_size()).to(device)
    
    # Tensorboard
    writer = SummaryWriter(config['experiment_name'])
    
    # 使用'# config'字典中的指定学习率和ε值设置优化器
    optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'], eps = 1e-9)
    
    # 初始化和全局步长变量
    initial_epoch = 0
    global_step = 0
    
    if config['preload']:
        model_filename = get_weights_file_path(config, config['preload'])
        print(f'Preloading model {model_filename}')
        state = torch.load(model_filename) # Loading model
        
        initial_epoch = state['epoch'] + 1
        optimizer.load_state_dict(state['optimizer_state_dict'])
        global_step = state['global_step']
        
    loss_fn = nn.CrossEntropyLoss(ignore_index = tokenizer_src.token_to_id('[PAD]'), label_smoothing = 0.1).to(device)

    for epoch in range(initial_epoch, config['num_epochs']):

        batch_iterator = tqdm(train_dataloader, desc = f'Processing epoch {epoch:02d}')

        for batch in batch_iterator:
            model.train() # Train the model
            encoder_input = batch['encoder_input'].to(device)
            decoder_input = batch['decoder_input'].to(device)
            encoder_mask = batch['encoder_mask'].to(device)
            decoder_mask = batch['decoder_mask'].to(device)
           
            encoder_output = model.encode(encoder_input, encoder_mask)
            decoder_output = model.decode(encoder_output, encoder_mask, decoder_input, decoder_mask)
            proj_output = model.project(decoder_output)

            label = batch['label'].to(device)

            loss = loss_fn(proj_output.view(-1, tokenizer_tgt.get_vocab_size()), label.view(-1))

            batch_iterator.set_postfix({f"loss": f"{loss.item():6.3f}"})
            
            writer.add_scalar('train loss', loss.item(), global_step)
            writer.flush()
            loss.backward()

            optimizer.step()

            optimizer.zero_grad()
            
            global_step += 1 
            
        run_validation(model, val_dataloader, tokenizer_src, tokenizer_tgt, config['seq_len'], device, lambda msg: batch_iterator.write(msg), global_step, writer)

        model_filename = get_weights_file_path(config, f'{epoch:02d}')
        torch.save({
            'epoch': epoch, # Current epoch
            'model_state_dict': model.state_dict(),# Current model state
            'optimizer_state_dict': optimizer.state_dict(), # Current optimizer state
            'global_step': global_step # Current global step 
        }, model_filename)

现在开始训练我们的模型！

if __name__ == '__main__':
    warnings.filterwarnings('ignore') # 忽略警告
    config = get_config() # 检索配置设置
    train_model(config) # 使用配置参数训练模型

结果如下：

Using device cuda
Downloading builder script:
6.08k/? [00:00<00:00, 391kB/s]
Downloading metadata:
161k/? [00:00<00:00, 11.0MB/s]
Downloading and preparing dataset opus_books/en-it (download: 3.14 MiB, generated: 8.58 MiB, post-processed: Unknown size, total: 11.72 MiB) to /root/.cache/huggingface/datasets/opus_books/en-it/1.0.0/e8f950a4f32dc39b7f9088908216cd2d7e21ac35f893d04d39eb594746af2daf...
Downloading data: 100%
3.30M/3.30M [00:00<00:00, 10.6MB/s]
Dataset opus_books downloaded and prepared to /root/.cache/huggingface/datasets/opus_books/en-it/1.0.0/e8f950a4f32dc39b7f9088908216cd2d7e21ac35f893d04d39eb594746af2daf. Subsequent calls will reuse this data.
Max length of source sentence: 309
Max length of target sentence: 274
....................................................................