Bert模型详解

import os
import random
import time

import numpy as np
import pandas as pd
import torch
from torch import nn
from torch.optim import Adam
from torch.utils.data import Dataset, DataLoader
from tqdm import tqdm
from transformers import BertModel
from transformers import BertTokenizer

train_data_path = '../data/Content_Store/006_game_first_tag_train_data.txt'
dev_data_path = '../data/Content_Store/006_game_first_tag_test_data.txt'
test_data_path = '../data/Content_Store/006_game_first_tag_test_data.txt'
label_path = '../data/Content_Store/first_tag_class.txt'

train_df = pd.read_csv(train_data_path, sep='##\$\$&&', header=None)
dev_df = pd.read_csv(dev_data_path, sep='##\$\$&&', header=None)
test_df = pd.read_csv(test_data_path, sep='##\$\$&&', header=None)

new_columns = ['text', 'label']
train_df = train_df.rename(columns=dict(zip(train_df.columns, new_columns)))
dev_df = dev_df.rename(columns=dict(zip(dev_df.columns, new_columns)))
test_df = test_df.rename(columns=dict(zip(test_df.columns, new_columns)))

real_labels = []
with open(label_path, 'r') as f:
    for row in f.readlines():
        real_labels.append(row.strip())

train_df.shape, dev_df.shape, real_labels

train_df.head(10)

BERT_PATH = '../modules/chinese-roberta-wwm-ext'
tokenizer = BertTokenizer.from_pretrained(BERT_PATH)

class MyDataset(Dataset):
    def __init__(self, df):
        self.texts = [
            tokenizer(
                text,
                padding='max_length',
                max_length=512,
                truncation=True,
                return_tensors="pt"
            )
            for text in df['text']
        ]
        self.labels = [label for label in df['label']]

    def __getitem__(self, idx):
        return self.texts[idx], self.labels[idx]

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

train_dataset = MyDataset(train_df)
dev_dataset = MyDataset(dev_df)
test_dataset = MyDataset(test_df)

class RoBertaClassifier(nn.Module):
    def __init__(self):
        super(RoBertaClassifier, self).__init__()
        self.bert = BertModel.from_pretrained(BERT_PATH)
        self.dropout = nn.Dropout(0.5)
        self.linear = nn.Linear(768, 8)
        self.relu = nn.ReLU()

    def forward(self, input_id, mask):
        _, pooled_output = self.bert(input_ids=input_id, attention_mask=mask, return_dict=False)
        dropout_output = self.dropout(pooled_output)
        linear_output = self.linear(dropout_output)
        final_layer = self.relu(linear_output)
        return final_layer


epoch = 5
batch_size = 32
lr = 2e-5
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
random_seed = 1999
save_path = './roberta_wwm_content_store_checkpoint'

def setup_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True

setup_seed(random_seed)

def save_model(save_name):
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    torch.save(model.state_dict(), os.path.join(save_path, save_name))
    ## model.save_pretrained(os.path.join(save_path, save_name))
    ## tokenizer.save_pretrained(os.path.join(save_path, save_name))

model = RoBertaClassifier()

if torch.cuda.device_count() > 1:
    print(f"使用 {torch.cuda.device_count()} 张 GPU!")
    model = nn.DataParallel(model)

criterion = nn.CrossEntropyLoss()
optimizer = Adam(model.parameters(), lr=lr)

model = model.to(device)
criterion = criterion.to(device)

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
dev_loader = DataLoader(dev_dataset, batch_size=batch_size)

best_dev_acc = 0
for epoch_num in range(epoch):
    total_acc_train = 0
    total_loss_train = 0
    for inputs, labels in tqdm(train_loader):
        ## 获取ids和mask，这些是transformer的输入
        input_ids = inputs['input_ids'].squeeze(1).to(device)
        masks = inputs['attention_mask'].to(device)
        labels = labels.to(device)
        ## 计算新的分类模型的输出值
        output = model(input_ids, masks)

        batch_loss = criterion(output, labels)
        batch_loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        acc = (output.argmax(dim=1) == labels).sum().item()
        total_acc_train += acc
        total_loss_train += batch_loss.item()

    ## 模型置为验证模式，不影响训练过程
    model.eval()
    total_acc_val = 0
    total_loss_val = 0
    with torch.no_grad():
        for inputs, labels in dev_loader:
            input_ids = inputs['input_ids'].squeeze(1).to(device)
            masks = inputs['attention_mask'].to(device)
            labels = labels.to(device)
            output = model(input_ids, masks)

            batch_loss = criterion(output, labels)
            acc = (output.argmax(dim=1) == labels).sum().item()
            total_acc_val += acc
            total_loss_val += batch_loss.item()

        print(f'''Epochs: {epoch_num + 1} 
          | Train Loss: {total_loss_train / len(train_dataset): .3f} 
          | Train Accuracy: {total_acc_train / len(train_dataset): .3f} 
          | Val Loss: {total_loss_val / len(dev_dataset): .3f} 
          | Val Accuracy: {total_acc_val / len(dev_dataset): .3f}''')

        if total_acc_val / len(dev_dataset) > best_dev_acc:
            best_dev_acc = total_acc_val / len(dev_dataset)
            ## save_model(model, tokenizer, 'best_' + str(best_dev_acc) + '_' + str(int(time.time())))
            save_model('best.pt')

    model.train()

save_model('last.pt')

model = RoBertaClassifier()
if torch.cuda.device_count() > 1:
    print(f"使用 {torch.cuda.device_count()} 张 GPU!")
    model = nn.DataParallel(model)
model.load_state_dict(torch.load(os.path.join(save_path, 'best.pt')))
model = model.to(device)
model.eval()


def evaluate(model, dataset):
    model.eval()
    test_loader = DataLoader(dataset, batch_size=128)
    total_acc_test = 0
    with torch.no_grad():
        for test_input, test_label in test_loader:
            input_id = test_input['input_ids'].squeeze(1).to(device)
            mask = test_input['attention_mask'].to(device)
            test_label = test_label.to(device)
            output = model(input_id, mask)
            acc = (output.argmax(dim=1) == test_label).sum().item()
            total_acc_test += acc
    print(f'Test Accuracy: {total_acc_test / len(dataset): .3f}')

evaluate(model, test_dataset)

while True:
    text = input('APP内容：')
    bert_input = tokenizer(
        text,
        padding='max_length',
        max_length=512,
        truncation=True,
        return_tensors="pt"
    )
    input_ids = bert_input['input_ids'].to(device)
    masks = bert_input['attention_mask'].unsqueeze(1).to(device)
    output = model(input_ids, masks)
    pred = output.argmax(dim=1)
    print("标签索引：" + str(pred))
    print("标签：" + str(real_labels[pred]))