Datawhale AI夏令营基于术语词典干预的机器翻译挑战比赛-task2(模型优化笔记)

 任务背景
目标是优化模型的训练和验证过程,并解决优化后出现的输入和目标批次大小不匹配的问题。

优化部分

数据清洗
参考baseline大佬笔记的部分做的数据清理
对缩写的单词进行了扩展,并且删除一些特殊字符,只保留一些标点符号和数字等。代码如下:

import contractions

def unicodeToAscii(text):
    return ''.join(c for c in unicodedata.normalize('NFD', text) if unicodedata.category(c) != 'Mn')

def preprocess_en(text):
    text = unicodeToAscii(text.strip())
    text = contractions.fix(text)
    text = re.sub(r'\([^)]*\)', '', text)
    text = re.sub(r"[^a-zA-Z0-9.!?]+", r" ", text)  # 保留数字
    return text

对(掌声)这种不该出现在翻译文本中的脏数据可以使用正则表达式剔除,代码如下:

def preprocess_zh(text):
    # 去除(掌声)这些脏数据
    text = re.sub(r'\([^)]*\)', '', text)
    text = re.sub(r"[^\u4e00-\u9fa5,。!?0-9]", "", text)  # 保留数字
    return text

效果挺好的,不然task3里跑完了会出现这种情况:

清理了以后就相当干净(虽然我也不知道为什么分数降了,主办方你说呢)

梯度累积

梯度累计(Gradient Accumulation)简介

梯度累计(Gradient Accumulation)是一种在深度学习训练过程中有效利用内存和计算资源的技术。它通过在多个小批次(micro-batches)上累积梯度,并在特定步数后进行一次参数更新,从而实现大批次训练的效果。这对于显存有限的硬件设备尤为重要。

背景和动机

在深度学习训练中,较大的批次大小通常能够带来更稳定的梯度估计和更快的收敛速度。然而,由于显存限制,我们往往无法直接使用大批次进行训练。梯度累计通过在多个小批次上累积梯度,实现大批次的训练效果,同时保持显存占用在可控范围内。

优势
  1. 降低显存需求

    • 梯度累计允许我们使用小批次进行训练,从而减少每次前向和反向传播所需的显存。
    • 在显存有限的环境下,能够训练更大规模的模型或使用更大的虚拟批次(accumulated batch)。
  2. 提高训练稳定性

    • 使用大批次通常能够提供更稳定的梯度估计,梯度累计模拟了这种效果,有助于提高训练过程的稳定性和收敛速度。
  3. 灵活性

    • 梯度累计可以与其他优化技术结合使用,例如混合精度训练,从而进一步提高训练效率和性能。
实现方法

实现梯度累计主要包括以下步骤:

  1. 定义累积步数

    • 指定在多少个小批次后进行一次梯度更新。
  2. 累积梯度

    • 在每个小批次上计算损失,并进行反向传播以累积梯度。
  3. 更新参数

    • 在累积步数达到指定次数后,进行一次参数更新,并清零梯度。

将梯度累积和验证步骤添加到现有的 train_model 函数中,按以下步骤修改了代码:

  1. 引入 gradient_accumulation_steps 参数,以控制梯度累积的步数。
  2. 修改训练循环,使其支持梯度累积。
  3. 添加每隔一定步数进行验证的逻辑。

以下是修改后的 train_model 函数:

def train_model(model, train_loader, dev_loader, optimizer, criterion, n_epochs, clip, gradient_accumulation_steps=1, validation_steps=100):
    steps = len(train_loader)  # 获取训练加载器中的总步骤数
    validation_steps = int(validation_steps * gradient_accumulation_steps)  # 计算验证步骤,每训练一定步数进行一次验证

    for epoch in range(n_epochs):  # 迭代训练
        model.train()  # 设置模型为训练模式
        epoch_loss = 0

        for step, (src, trg) in enumerate(train_loader, 1):  # 遍历训练数据
            src, trg = src.to(device), trg.to(device)  # 移动数据到设备
            optimizer.zero_grad()  # 梯度清零
            output = model(src, trg)  # 前向传播
            output_dim = output.shape[-1]
            output = output[1:].view(-1, output_dim)  # 调整输出形状
            trg = trg[1:].view(-1)
            loss = criterion(output, trg)  # 计算损失

            if gradient_accumulation_steps > 1:
                loss = loss / gradient_accumulation_steps  # 如果需要累积梯度,损失除以累积步数

            loss.backward()  # 反向传播

            if step % gradient_accumulation_steps == 0 or step == steps:
                clip_grad_norm_(model.parameters(), clip)  # 裁剪梯度
                optimizer.step()  # 更新参数
                model.zero_grad()  # 清零模型梯度

            epoch_loss += loss.item()  # 累加损失

            if step % validation_steps == 0:
                model.eval()  # 设置模型为评估模式
                validation_loss = 0
                with torch.no_grad():
                    for src_val, trg_val in dev_loader:  # 遍历验证数据
                        src_val, trg_val = src_val.to(device), trg_val.to(device)  # 移动数据到设备
                        output_val = model(src_val, trg_val)  # 前向传播
                        output_dim_val = output_val.shape[-1]
                        output_val = output_val[1:].view(-1, output_dim_val)  # 调整输出形状
                        trg_val = trg_val[1:].view(-1)
                        val_loss = criterion(output_val, trg_val)  # 计算损失
                        validation_loss += val_loss.item()  # 累加验证损失
                validation_loss /= len(dev_loader)  # 计算平均验证损失
                print(f'Validation Loss: {validation_loss}')  # 打印验证损失

        print(f'Epoch {epoch+1}/{n_epochs} Loss: {epoch_loss/len(train_loader)}')  # 打印每个epoch的损失

混合精度训练

混合精度训练简介

混合精度训练(Mixed Precision Training)是一种利用半精度(16位浮点数,FP16)和单精度(32位浮点数,FP32)结合进行神经网络训练的技术。这种方法能够显著提高训练速度和减少显存占用,同时还能保持模型的精度和稳定性。

背景和动机

传统的深度学习模型训练通常使用单精度浮点数(FP32),但随着模型的复杂性和规模不断增加,训练时间和显存需求也随之增长。混合精度训练通过在不影响模型精度的前提下使用半精度浮点数,可以有效地减少计算资源的消耗。

优势
  1. 提高训练速度

    • 使用半精度浮点数可以显著提高GPU的吞吐量,因为FP16运算速度通常是FP32的两倍。
    • 通过更快的计算,训练过程可以大大加速。
  2. 减少显存占用

    • FP16占用的显存仅为FP32的一半,这意味着可以在同样的显存空间内训练更大规模的模型或使用更大的批次(batch size)。
    • 这对于显存有限的GPU尤其有利,可以有效缓解显存瓶颈问题。
  3. 保持精度和稳定性

    • 混合精度训练通过动态缩放(loss scaling)等技术,能够在保证数值稳定性的同时,避免梯度下溢问题。
    • 结合FP16和FP32的优点,既能享受半精度浮点数带来的效率提升,又能保留单精度浮点数的精度优势。
实现方法

实现混合精度训练主要包括以下几个步骤:

  1. 模型和优化器的初始化

    • 初始化模型和优化器,与传统训练方式相同。
  2. 使用 autocast 进行前向传播

    • 使用 torch.cuda.amp.autocast 上下文管理器进行前向传播计算。
    • 在该上下文中,所有计算将默认使用半精度浮点数进行,除非某些操作需要单精度浮点数。
  3. 损失缩放和反向传播

    • 使用 torch.cuda.amp.GradScaler 对损失进行缩放,然后执行反向传播。
    • 通过损失缩放,可以避免梯度下溢问题,并确保梯度计算的数值稳定性。
  4. 参数更新和缩放器更新

    • 使用缩放后的梯度进行优化器的参数更新。
    • 在更新参数后,调用缩放器的 update 方法以调整下次训练步骤中的缩放系数。

为了将混合精度训练与上述代码结合,按以下步骤修改了代码:

  1. 引入 torch.cuda.amp 模块。
  2. 在训练循环中使用 autocast 上下文管理器进行前向传播。
  3. 使用 GradScaler 进行梯度缩放和反向传播。

以下是修改后的 train_model 函数:

from torch.cuda.amp import autocast, GradScaler  # 导入自动混合精度训练所需的autocast和GradScaler模块

def train_model(model, train_loader, dev_loader, optimizer, criterion, n_epochs, clip, gradient_accumulation_steps=1, validation_steps=100):
    steps = len(train_loader)  # 获取训练加载器中的总步骤数
    validation_steps = int(validation_steps * gradient_accumulation_steps)  # 计算验证步骤,每训练一定步数进行一次验证
    scaler = GradScaler()  # 初始化GradScaler对象,用于缩放梯度

    for epoch in range(n_epochs):  # 迭代训练
        model.train()  # 设置模型为训练模式
        epoch_loss = 0

        for step, (src, trg) in enumerate(train_loader, 1):  # 遍历训练数据
            src, trg = src.to(device), trg.to(device)  # 移动数据到设备
            optimizer.zero_grad()  # 梯度清零
            
            # 前向传播使用autocast上下文管理器
            with autocast(enabled=True):  # 在autocast上下文管理器中进行前向传播,启用混合精度
                output = model(src, trg)  # 前向传播,计算模型输出
                output_dim = output.shape[-1]
                output = output[1:].view(-1, output_dim)  # 调整输出形状
                trg = trg[1:].view(-1)
                loss = criterion(output, trg)  # 计算损失

                if gradient_accumulation_steps > 1:
                    loss = loss / gradient_accumulation_steps  # 如果需要累积梯度,损失除以累积步数

            # 使用scaler缩放损失并进行反向传播
            scaler.scale(loss).backward()  # 缩放损失并进行反向传播

            if step % gradient_accumulation_steps == 0 or step == steps:
                scaler.unscale_(optimizer)  # 在进行梯度裁剪前,取消缩放优化器参数
                clip_grad_norm_(model.parameters(), clip)  # 裁剪梯度,防止梯度爆炸
                scaler.step(optimizer)  # 使用缩放后的梯度进行优化步骤,更新模型参数
                scaler.update()  # 更新缩放器,调整下次的缩放系数
                model.zero_grad()  # 清零模型梯度

            epoch_loss += loss.item()  # 累加损失

            if step % validation_steps == 0:
                model.eval()  # 设置模型为评估模式
                validation_loss = 0
                with torch.no_grad():
                    for src_val, trg_val in dev_loader:  # 遍历验证数据
                        src_val, trg_val = src_val.to(device), trg_val.to(device)  # 移动数据到设备
                        with autocast(enabled=True):  # 在autocast上下文管理器中进行前向传播,启用混合精度
                            output_val = model(src_val, trg_val)  # 前向传播
                            output_dim_val = output_val.shape[-1]
                            output_val = output_val[1:].view(-1, output_dim_val)  # 调整输出形状
                            trg_val = trg_val[1:].view(-1)
                            val_loss = criterion(output_val, trg_val)  # 计算损失
                        validation_loss += val_loss.item()  # 累加验证损失
                validation_loss /= len(dev_loader)  # 计算平均验证损失
                print(f'Validation Loss: {validation_loss}')  # 打印验证损失

        print(f'Epoch {epoch+1}/{n_epochs} Loss: {epoch_loss/len(train_loader)}')  # 打印每个epoch的损失

 问题描述
在梯度累计和混合精度计算最初实现中,我们遇到了输入批次大小与目标批次大小不匹配的错误。这种错误通常发生在处理变长序列时,输入和目标序列的长度不一致。

错误信息:

ValueError: Expected input batch_size (1376) to match target batch_size (1472).
 

 问题分析
此问题主要出现在以下几个方面:
1. 序列长度不一致:输入序列和目标序列的填充方式可能不一致,导致它们在计算损失时形状不匹配。
2. 模型输出和目标序列的调整:在计算损失之前,模型输出和目标序列需要调整为相同的形状。

 优化步骤
为了解决上述问题,我们对代码进行了以下优化:

 1. 调整数据加载和填充方式
我们使用 `pad_sequence` 函数对输入和目标序列进行填充,确保它们在一个批次内具有相同的长度。

def collate_fn(batch):
    src_batch, trg_batch = zip(*batch)
    src_batch = pad_sequence([torch.tensor(x) for x in src_batch], padding_value=0, batch_first=True)  # 对输入进行填充
    trg_batch = pad_sequence([torch.tensor(x) for x in trg_batch], padding_value=0, batch_first=True)  # 对目标进行填充
    return src_batch, trg_batch


 

 2. 调整模型输出和目标序列的形状
在训练和验证过程中,我们在计算损失之前调整模型输出和目标序列的形状,使其匹配。

 

def train_model(model, train_loader, dev_loader, optimizer, criterion, n_epochs, clip, gradient_accumulation_steps=1, validation_steps=100):
    steps = len(train_loader)
    validation_steps = int(validation_steps * gradient_accumulation_steps)
    scaler = GradScaler()

    for epoch in range(n_epochs):
        model.train()
        epoch_loss = 0

        for step, (src, trg) in enumerate(train_loader, 1):
            src, trg = src.to(device), trg.to(device)
            optimizer.zero_grad()

            with autocast(enabled=True):
                output = model(src, trg)
                output_dim = output.shape[-1]
                output = output[:, 1:].contiguous().view(-1, output_dim)  # 调整输出形状
                trg = trg[:, 1:].contiguous().view(-1)  # 调整目标形状
                loss = criterion(output, trg)

                if gradient_accumulation_steps > 1:
                    loss = loss / gradient_accumulation_steps

            scaler.scale(loss).backward()

            if step % gradient_accumulation_steps == 0 or step == steps:
                scaler.unscale_(optimizer)
                clip_grad_norm_(model.parameters(), clip)
                scaler.step(optimizer)
                scaler.update()
                model.zero_grad()

            epoch_loss += loss.item()

            if step % validation_steps == 0:
                model.eval()
                validation_loss = 0
                with torch.no_grad():
                    for src_val, trg_val in dev_loader:
                        src_val, trg_val = src_val.to(device), trg_val.to(device)
                        with autocast(enabled=True):
                            output_val = model(src_val, trg_val)
                            output_dim_val = output_val.shape[-1]
                            output_val = output_val[:, 1:].contiguous().view(-1, output_dim_val)  # 调整输出形状
                            trg_val = trg_val[:, 1:].contiguous().view(-1)  # 调整目标形状
                            val_loss = criterion(output_val, trg_val)
                        validation_loss += val_loss.item()
                validation_loss /= len(dev_loader)
                print(f'Validation Loss: {validation_loss}')

        print(f'Epoch {epoch+1}/{n_epochs} Loss: {epoch_loss/len(train_loader)}')

 3. 处理训练和验证数据
确保在处理训练和验证数据时,输入和目标序列保持一致,并在计算损失之前正确调整形状。


# 数据加载函数
def load_data(train_path, dev_en_path, dev_zh_path, test_en_path):
    en_train_data = read_data(train_path)
    zh_train_data = read_data(train_path)
    en_dev_data = read_data(dev_en_path)
    zh_dev_data = read_data(dev_zh_path)
    en_test_data = read_data(test_en_path)
    
    train_data = preprocess_data(en_train_data, zh_train_data)
    dev_data = preprocess_data(en_dev_data, zh_dev_data)
    test_data = preprocess_data(en_test_data, [])
    
    en_vocab, zh_vocab = build_vocab(train_data)
    
    train_dataset = TranslationDataset(train_data, en_vocab, zh_vocab)
    dev_dataset = TranslationDataset(dev_data, en_vocab, zh_vocab)
    test_dataset = TranslationDataset(test_data, en_vocab, zh_vocab)
    
    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_fn)
    dev_loader = DataLoader(dev_dataset, batch_size=BATCH_SIZE, shuffle=False, collate_fn=collate_fn)
    test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, collate_fn=collate_fn)
    
    return train_loader, dev_loader, test_loader, en_vocab, zh_vocab

 4. 完整的优化代码
优化后的代码结合了以上所有改进,确保输入和目标序列的形状一致,解决了输入批次大小与目标批次大小不匹配的问题。


import os
import re
import contractions
import unicodedata
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.nn.utils import clip_grad_norm_
from torchtext.data.metrics import bleu_score
from torch.utils.data import Dataset, DataLoader
from torchtext.data.utils import get_tokenizer
from torchtext.vocab import build_vocab_from_iterator
from typing import List, Tuple
import jieba
import random
from torch.nn.utils.rnn import pad_sequence
import sacrebleu
import time
import math
import bitsandbytes as bnb
from torch.cuda.amp import autocast, GradScaler
from torch.utils.checkpoint import checkpoint

# 检测GPU是否可用
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 定义tokenizer
en_tokenizer = get_tokenizer('spacy', language='en_core_web_trf')
zh_tokenizer = lambda x: list(jieba.cut(x))

# 数据清洗函数
def unicode_to_ascii(text):
    return ''.join(c for c in unicodedata.normalize('NFD', text) if unicodedata.category(c) != 'Mn')

def preprocess_en(text):
    text = unicode_to_ascii(text.strip())
    text = contractions.fix(text)
    text = re.sub(r'\([^)]*\)', '', text)
    text = re.sub(r"[^a-zA-Z0-9.!?]+", r" ", text)
    return text

def preprocess_zh(text):
    text = re.sub(r'\([^)]*\)', '', text)
    text = re.sub(r"[^\u4e00-\u9fa5,。!?0-9]", "", text)
    return text

# 读取数据函数
def read_data(file_path: str) -> List[str]:
    with open(file_path, 'r', encoding='utf-8') as f:
        return [line.strip() for line in f]

# 数据预处理函数
def preprocess_data(en_data: List[str], zh_data: List[str]) -> List[Tuple[List[str], List[str]]]:
    processed_data = []
    for en, zh in zip(en_data, zh_data):
        en = preprocess_en(en)
        zh = preprocess_zh(zh)
        en_tokens = en_tokenizer(en.lower())[:MAX_LENGTH]
        zh_tokens = zh_tokenizer(zh)[:MAX_LENGTH]
        if en_tokens and zh_tokens:
            processed_data.append((en_tokens, zh_tokens))
    return processed_data

# 构建词汇表
def build_vocab(data: List[Tuple[List[str], List[str]]]):
    en_vocab = build_vocab_from_iterator(
        (en for en, _ in data),
        specials=['<unk>', '<pad>', '<bos>', '<eos>']
    )
    zh_vocab = build_vocab_from_iterator(
        (zh for _, zh in data),
        specials=['<unk>', '<pad>', '<bos>', '<eos>']
    )
    en_vocab.set_default_index(en_vocab['<unk>'])
    zh_vocab.set_default_index(zh_vocab['<unk>'])
    return en_vocab, zh_vocab

class TranslationDataset(Dataset):
    def __init__(self, data: List[Tuple[List[str], List[str]]], en_vocab, zh_vocab):
        self.data = data
        self.en_vocab = en_vocab
        self.zh_vocab = zh_vocab

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        en, zh = self.data[idx]
        en_indices = [self```python
        en_indices = [self.en_vocab['<bos>']] + [self.en_vocab[token] for token in en] + [self.en_vocab['<eos>']]
        zh_indices = [self.zh_vocab['<bos>']] + [self.zh_vocab[token] for token in zh] + [self.zh_vocab['<eos>']]
        return en_indices, zh_indices

def collate_fn(batch):
    en_batch, zh_batch = [], []
    for en_item, zh_item in batch:
        if en_item and zh_item:
            en_batch.append(torch.tensor(en_item))
            zh_batch.append(torch.tensor(zh_item))
    
    if not en_batch or not zh_batch:
        return torch.tensor([]), torch.tensor([])
    
    en_batch = nn.utils.rnn.pad_sequence(en_batch, batch_first=True, padding_value=en_vocab['<pad>'])
    zh_batch = nn.utils.rnn.pad_sequence(zh_batch, batch_first=True, padding_value=zh_vocab['<pad>'])
    
    return en_batch, zh_batch

# 数据加载函数
def load_data(train_path: str, dev_en_path: str, dev_zh_path: str, test_en_path: str):
    en_train_data = read_data(train_path)
    zh_train_data = read_data(train_path)
    en_dev_data = read_data(dev_en_path)
    zh_dev_data = read_data(dev_zh_path)
    en_test_data = read_data(test_en_path)
    
    train_data = preprocess_data(en_train_data, zh_train_data)
    dev_data = preprocess_data(en_dev_data, zh_dev_data)
    test_data = preprocess_data(en_test_data, [])
    
    en_vocab, zh_vocab = build_vocab(train_data)
    
    train_dataset = TranslationDataset(train_data, en_vocab, zh_vocab)
    dev_dataset = TranslationDataset(dev_data, en_vocab, zh_vocab)
    test_dataset = TranslationDataset(test_data, en_vocab, zh_vocab)
    
    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_fn)
    dev_loader = DataLoader(dev_dataset, batch_size=BATCH_SIZE, shuffle=False, collate_fn=collate_fn)
    test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, collate_fn=collate_fn)
    
    return train_loader, dev_loader, test_loader, en_vocab, zh_vocab

class Encoder(nn.Module):
    def __init__(self, input_dim, emb_dim, hid_dim, n_layers, dropout):
        super().__init__()
        self.embedding = nn.Embedding(input_dim, emb_dim)
        self.rnn = nn.GRU(emb_dim, hid_dim, n_layers, dropout=dropout, batch_first=True)
        self.dropout = nn.Dropout(dropout)

    def forward(self, src):
        embedded = self.dropout(self.embedding(src))
        outputs, hidden = self.rnn(embedded)
        return outputs, hidden

class Decoder(nn.Module):
    def __init__(self, output_dim, emb_dim, hid_dim, n_layers, dropout):
        super().__init__()
        self.output_dim = output_dim
        self.embedding = nn.Embedding(output_dim, emb_dim)
        self.rnn = nn.GRU(emb_dim, hid_dim, n_layers, dropout=dropout, batch_first=True)
        self.fc_out = nn.Linear(hid_dim, output_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, input, hidden):
        embedded = self.dropout(self.embedding(input))
        output, hidden = self.rnn(embedded, hidden)
        prediction = self.fc_out(output.squeeze(1))
        return prediction, hidden

class Seq2Seq(nn.Module):
    def __init__(self, encoder, decoder, device):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.device = device

    def forward(self, src, trg, teacher_forcing_ratio=0.5):
        batch_size = src.shape[0]
        trg_len = trg.shape[1]
        trg_vocab_size = self.decoder.output_dim
        outputs = torch.zeros(batch_size, trg_len, trg_vocab_size).to(self.device)
        _, hidden = self.encoder(src)
        input = trg[:, 0].unsqueeze(1)

        for t in range(1, trg_len):
            output, hidden = self.decoder(input, hidden)
            outputs[:, t, :] = output
            teacher_force = random.random() < teacher_forcing_ratio
            top1 = output.argmax(1)
            input = trg[:, t].unsqueeze(1) if teacher_force else top1.unsqueeze(1)

        return outputs

def initialize_model(input_dim, output_dim, emb_dim, hid_dim, n_layers, dropout, device):
    enc = Encoder(input_dim, emb_dim, hid_dim, n_layers, dropout)
    dec = Decoder(output_dim, emb_dim, hid_dim, n_layers, dropout)
    model = Seq2Seq(enc, dec, device).to(device)
    return model

def initialize_optimizer(model, learning_rate=0.001):
    return bnb.optim.AdamW(params=model.parameters(), lr=learning_rate, weight_decay=0.0, optim_bits=8)

def epoch_time(start_time, end_time):
    elapsed_time = end_time - start_time
    elapsed_mins = int(elapsed_time / 60)
    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))
    return elapsed_mins, elapsed_secs

def train(model, iterator, optimizer, criterion, clip, gradient_accumulation_steps=1):
    model.train()
    epoch_loss = 0
    scaler = GradScaler()

    for i, batch in enumerate(iterator):
        src, trg = batch
        if src.numel() == 0 or trg.numel() == 0:
            continue
        
        src, trg = src.to(DEVICE), trg.to(DEVICE)
        
        optimizer.zero_grad()
        
        with autocast():
            output = model(src, trg)
            output_dim = output.shape[-1]
            output = output[:, 1:].contiguous().view(-1, output_dim)
            trg = trg[:, 1:].contiguous().view(-1)
            loss = criterion(output, trg) / gradient_accumulation_steps

        scaler.scale(loss).backward()

        if (i + 1) % gradient_accumulation_steps == 0:
            scaler.unscale_(optimizer)
            clip_grad_norm_(model.parameters(), clip)
            scaler.step(optimizer)
            scaler.update()
            model.zero_grad()
        
        epoch_loss += loss.item()

    return epoch_loss / len(iterator)

def evaluate(model, iterator, criterion):
    model.eval()
    epoch_loss = 0
    with torch.no_grad():
        for i, batch in enumerate(iterator):
            src, trg = batch
            if src.numel() == 0 or trg.numel() == 0:
                continue
            
            src, trg = src.to(DEVICE), trg.to(DEVICE)
            
            with autocast():
                output = model(src, trg, 0)
                output_dim = output.shape[-1]
                output = output[:, 1:].contiguous().view(-1, output_dim)
                trg = trg[:, 1:].contiguous().view(-1)
                loss = criterion(output, trg)
                epoch_loss += loss.item()
        
    return epoch_loss / len(iterator)

def translate_sentence(sentence, src_vocab, trg_vocab, model, device, max_length=50):
    model.eval()
    if isinstance(sentence, str):
        tokens = [token for token in en_tokenizer(sentence)]
    else:
        tokens = [str(token) for token in sentence]
        
    tokens = ['<bos>'] + tokens + ['<eos>']
    src_indexes = [src_vocab[token] for token in tokens]
    src_tensor = torch.LongTensor(src_indexes).unsqueeze(0).to(device)
    
    with torch.no_grad():
        encoder_outputs, hidden = model.encoder(src_tensor)
    
    trg_indexes = [trg_vocab['<bos>']]
    for i in range(max_length):
        trg_tensor = torch.LongTensor([trg_indexes[-1]]).to(device)
        
        with torch.no_grad():
            output, hidden = model.decoder(trg_tensor, hidden, encoder_outputs)
        
        pred_token = output.argmax(1).item()
        trg_indexes.append(pred_token)
        
        if pred_token == trg_vocab['<eos>']:
            break
    
    trg_tokens = [trg_vocab.get_itos()[i] for i in trg_indexes]
    return trg_tokens[1:-1]

def calculate_bleu(dev_loader, src_vocab, trg_vocab, model, device):
    translated_sentences = []
    references = []

    for src, trg in dev_loader:
        src = src.to(device)
        translation = translate_sentence(src, src_vocab, trg_vocab, model, device)
        translated_sentences.append(' '.join(translation))

        for t in trg:
            ref_str = ' '.join([trg_vocab.get_itos()[idx] for idx in t.tolist() if idx not in [trg_vocab['<bos>'], trg_vocab['<eos>'], trg_vocab['<pad>']]])
            references.append(ref_str)

    bleu = sacrebleu.corpus_bleu(translated_sentences, [references])
    return bleu.score

def train_model(model, train_iterator, valid_iterator, optimizer, criterion, N_EPOCHS=10, CLIP=1, gradient_accumulation_steps=1):
    best_valid_loss = float('inf')
    
    for epoch in range(N_EPOCHS):
        start_time = time.time()
        
        print(f"Starting Epoch {epoch + 1}")
        train_loss = train(model, train_iterator, optimizer, criterion, CLIP, gradient_accumulation_steps)
        valid_loss = evaluate(model, valid_iterator, criterion)
        
        end_time = time.time()
        epoch_mins, epoch_secs = epoch_time(start_time, end_time)
        
        if valid_loss < best_valid_loss:
            best_valid_loss = valid_loss
            if not os.path.exists('../model'):
                os.makedirs('../model')
            torch.save(model.state_dict(), '../model/best-model_test.pt')
        
        print(f'Epoch: {epoch+1:02} | Time: {epoch_mins}m {epoch_secs}s')
        print(f'\tTrain Loss: {train_loss:.3f} | Train PPL: {math.exp(train_loss):7.3f}')
        print(f'\t Val. Loss: {valid_loss:.3f} |  Val. PPL: {math.exp(valid_loss):7.3f}')

# 定义常量
MAX_LENGTH = 1000  # 最大句子长度
BATCH_SIZE = 32
DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
N = 100   # 采样训练集的数量

train_path = '../dataset/train.txt'
dev_en_path = '../dataset/dev_en.txt'
dev_zh_path = '../dataset/dev_zh.txt'
test_en_path = '../dataset/test_en.txt'

train_loader, dev_loader, test_loader, en_vocab, zh_vocab = load_data(
    train_path, dev_en_path, dev_zh_path, test_en_path
)

print(f"英语词汇表大小: {len(en_vocab)}")
print(f"中文词汇表大小: {len(zh_vocab)}")
print(f"训练集大小: {len(train_loader.dataset)}")
print(f"开发集大小: {len(dev_loader.dataset)}")
print(f"测试集大小: {len(test_loader.dataset)}")

# 主函数
if __name__ == '__main__':
    N_EPOCHS = 10
    CLIP = 1
    gradient_accumulation_steps = 2  # 根据需要调整
    # 模型参数
    INPUT_DIM = len(en_vocab)
    OUTPUT_DIM = len(zh_vocab)
    EMB_DIM = 128
    HID_DIM = 256
    N_LAYERS = 2  # 调整层数
    DROPOUT = 0.5
    
    # 初始化模型
    model = initialize_model(INPUT_DIM, OUTPUT_DIM, EMB_DIM, HID_DIM, N_LAYERS, DROPOUT, DEVICE)
    print(f'The model has {sum(p.numel() for p in model.parameters() if p.requires_grad):,} trainable parameters')

    # 定义损失函数
    criterion = nn.CrossEntropyLoss(ignore_index=zh_vocab['<pad>'])
    # 初始化优化器
    optimizer = initialize_optimizer(model)

    # 训练模型
    train_model(model, train_loader, dev_loader, optimizer, criterion, N_EPOCHS, CLIP, gradient_accumulation_steps)
    
    # 加载最佳模型
    model.load_state_dict(torch.load('../model/best-model_test.pt'))

    # 计算BLEU分数
    bleu_score = calculate_bleu(dev_loader, en_vocab, zh_vocab, model, DEVICE)
    print(f'BLEU score = {bleu_score:.2f}')

    # 确保结果目录存在
    if not os.path.exists('../results'):
        os.makedirs('../results')

    with open('../results/submit_test.txt', 'w') as f:
        translated_sentences = []
        for batch in test_loader:
            src, _ = batch
            src = src.to(DEVICE)
            translated = translate_sentence(src[0], en_vocab, zh_vocab, model, DEVICE)
            results = "".join(translated)
            f.write(results + '\n')


 

通过这些优化,可以在有限的GPU资源下高效地训练大模型,并减轻显存压力。

 优化总结

主要进行了以下优化:
1. 数据清洗:处理输入中的特殊字符和缩写,确保模型输入的一致性。
2. 数据加载和填充:使用 `pad_sequence` 对输入和目标序列进行填充,确保它们在一个批次内具有相同的长度。
3. 模型结构:构建了基于 GRU 的 Seq2Seq 模型,并在编码器和解码器中加入了合适的层数和丢弃率。
4. 训练优化:使用自动混合精度(AMP)和梯度累积来减少显存占用,提高训练速度。
5. 评估和推理:通过计算 BLEU 分数评估模型性能,并实现了模型推理和结果保存。

但是task2改了跑起来还是一言难尽。

更神奇的是性能确确实实的释放出来了,但是就是拉不满。。。

改了很多版,熬了三天凌晨,只能说效果不尽人意。。。(悲)希望task3能好点

  • 11
    点赞
  • 8
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值