NLP-分类-BertMLP

NLP-分类-BertMLP

• 数据集
• 代码

# -*- ecoding: utf-8 -*-
# @Author: SuperLong
# @Email: miu_zxl@163.com
# @Time: 2024/8/30 8:54
import os
import time
import torch
import argparse
from pathlib import Path
from torch.utils.data import DataLoader,Dataset
from torch import nn as nn
from torch.optim import AdamW
from sklearn.metrics import accuracy_score
from transformers import BertTokenizer,BertModel
import warnings
warnings.filterwarnings('ignore')
path = Path(__file__).parent

def load_data(args,stype_str):
    text, label, max_len = [], [], []
    with open(args.data_path + f"/{stype_str}.txt", "r", encoding='utf-8') as f:
        for line in f.readlines():
            if not line:
                continue
            text_, label_ = line.strip().split('\t')
            text.append(text_)
            label.append(label_)
            max_len.append(len(text_))
    if stype_str == 'train':
        return text, label, max(max_len)
    else:
        return text, label


class Datasets(Dataset):
    def __init__(self, text, label, max_len):
        super(Datasets, self).__init__()
        self.text = text
        self.label = label
        self.max_len = max_len
        self.tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')

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

    def __getitem__(self, idx):
        text = self.text[idx][:self.max_len]
        label = self.label[idx]

        text_f = self.tokenizer.tokenize(text)
        text_f = ['[CLS]'] + text_f #+ ['[SEP]']
        text_id = self.tokenizer.convert_tokens_to_ids(text_f)
        mask = [1] * len(text_f) + [0] * (self.max_len + 2 - len(text_f))
        text_id += [0] * (self.max_len + 2 - len(text_f))
        label_id = int(label)

        text_id = torch.tensor(text_id,)
        label_id = torch.tensor(label_id,)
        mask = torch.tensor(mask,)

        return (text_id, mask), label_id


class MyBertClass(nn.Module):
    def __init__(self,class_num,fiff_dim,droupout=0.1):
        super(MyBertClass,self).__init__()
        self.dropout = nn.Dropout(droupout)
        self.bert = BertModel.from_pretrained('bert-base-chinese')
        self.linear = nn.Linear(fiff_dim, class_num)
        # self.device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
        self.device = torch.device('cuda:0')
        for param in self.bert.parameters():
            param.requires_grid = True

    def forward(self,x):
        token_ids,mask = x[0].to(self.device),x[1].to(self.device)
        hidden_out = self.bert(token_ids,attention_mask=mask,output_hidden_states=False) # (bitch_size,content_len,bem_dim)
        # y = self.dropout(y.pooler_output)
        return self.linear(hidden_out.pooler_output)

def main(args):
    start = time.time()
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    train_text, train_label, max_len = load_data(args,'train')
    train_data = Datasets(train_text, train_label, max_len)
    train_loader = DataLoader(train_data, batch_size=args.batch_size, shuffle=True)

    dev_text, dev_label = load_data(args,'dev')
    dev_data = Datasets(dev_text, dev_label,max_len)
    dev_loader = DataLoader(dev_data, batch_size=args.batch_size, shuffle=True)
    model = MyBertClass(args.class_num, args.linear_dim)
    model.to(device)
    opt = AdamW(model.parameters(), lr=args.learning_rate)
    loss_f = nn.CrossEntropyLoss()

    best_acc = float("-inf")
    for epoch in range(args.epochs):
        model.train()
        loss_sim, count = 0, 0
        for step, (batch_text, batch_label) in enumerate(train_loader):
            batch_label = batch_label.to(device)
            opt.zero_grad()
            out = model(batch_text)
            loss = loss_f(out, batch_label)
            loss.backward()
            opt.step()
            loss_sim += loss
            count += 1
            if step % 10 == 0:
                meg = 'Epoch:{0} Step:{1} Loss:{2:.4f}'
                print(meg.format(epoch, step, loss_sim / count))
                loss_sim, count = 0, 0

        model.eval()
        pred_list, true_list = [], []
        with torch.no_grad():
            for batch_text, batch_label in dev_loader:
                batch_label.to(device)
                y_pred = model(batch_text)
                y_pred = torch.argmax(y_pred,dim=1).cpu().numpy().tolist()
                pred_list.append(y_pred)
                batch_label = batch_label.cpu().numpy().tolist()
                true_list.append(batch_label)
        acc = accuracy_score(pred_list, true_list)

        if acc < best_acc:
            best_acc = acc
            torch.save(model.state_dict(), 'best_model-BertMLP.pth')
        print("测试结果为：{:.4f}".format(acc))
    cost_time = time.time() - start
    print(f'训练结束，共耗时：{cost_time:.2f}')


if __name__ == '__main__':
    args = argparse.ArgumentParser()
    args.add_argument('--epochs',type=int,default=1)
    args.add_argument('--batch_size',type=int,default=64)
    args.add_argument('--learning_rate',type=float,default=0.0001)
    args.add_argument('--class_num',type=int,default=10)
    args.add_argument('--linear_dim',type=int,default=768)
    args.add_argument("--data_path", type=str, default=os.path.join(path, "data", "class_data"))
    args.add_argument("--model_path", type=str, default=os.path.join(path, "bset_model-BertMLP.pth"))
    args = args.parse_args()
    main(args)

训练损失不收敛有跳动，请指正！

目前仅专注于NLP的技术学习和分享
感谢大家的关注与支持！