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)
训练损失不收敛有跳动,请指正!