From 9cff4423ed76c56b9d082cf22599c8100ea612a0 Mon Sep 17 00:00:00 2001 From: Link_pursuit <1149830513@qq.com> Date: Wed, 27 Mar 2024 11:30:04 +0800 Subject: [PATCH] Update bert_train.py --- BERT/bert_train.py | 156 +++++++++++++++++++++++++++++++-------------- 1 file changed, 107 insertions(+), 49 deletions(-) diff --git a/BERT/bert_train.py b/BERT/bert_train.py index 50d6a93..724d716 100644 --- a/BERT/bert_train.py +++ b/BERT/bert_train.py @@ -3,27 +3,25 @@ import numpy as np from sklearn.model_selection import train_test_split from transformers import BertTokenizer, BertForSequenceClassification, Trainer, TrainingArguments import torch -from sklearn.metrics import f1_score, roc_auc_score, accuracy_score, precision_score, recall_score, average_precision_score, balanced_accuracy_score +from sklearn.metrics import f1_score, accuracy_score, precision_score, recall_score +import logging +import os +import data_split -df = pd.read_csv("../Dataset/autoirt_shuffled.csv") +train_df, val_df, test_df = data_split.split_dataset("../Dataset/new_autoirt.csv") -# 首先将数据分为训练集和测试集 -train_df, test_df = train_test_split(df, test_size=0.2, random_state=42) +# 打印数据集大小 +print(len(train_df), len(test_df), len(val_df)) -# 计算剩余数据的20%（相当于总数据的10%）作为验证集 -val_size = 0.2 * len(df) / (len(train_df) + len(test_df)) -train_df, val_df = train_test_split(train_df, test_size=val_size, random_state=42) +# 加载tokenizer +tokenizer = BertTokenizer.from_pretrained('bert-base-uncased') -# train_df 现在是训练集，val_df 是验证集，test_df 是测试集 -print(len(train_df),len(test_df),len(val_df)) +# 编码数据 +train_encodings = tokenizer(list(train_df['element']), truncation=True, padding=True) +val_encodings = tokenizer(list(val_df['element']), truncation=True, padding=True) +test_encodings = tokenizer(list(test_df['element']), truncation=True, padding=True) -tokenizers = BertTokenizer.from_pretrained('bert-base-uncased') - -train_encodings = tokenizers(list(train_df['element']), truncation=True, padding=True) -val_encodings = tokenizers(list(val_df['element']), truncation=True, padding=True) -test_encodings = tokenizers(list(test_df['element']), truncation=True, padding=True) - -# 将编码后的数据转换为Dataset格式 +# 定义数据集类 class CommentsDataset(torch.utils.data.Dataset): def __init__(self, encodings, labels): self.encodings = encodings @@ -37,30 +35,31 @@ class CommentsDataset(torch.utils.data.Dataset): def __len__(self): return len(self.labels) - +# 创建数据集实例 train_dataset = CommentsDataset(train_encodings, list(train_df['label'])) val_dataset = CommentsDataset(val_encodings, list(val_df['label'])) test_dataset = CommentsDataset(test_encodings, list(test_df['label'])) -# 加载BertForSequenceClassification模型 +# 加载模型 model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=6) # 定义训练参数 training_args = TrainingArguments( - output_dir='./BERT_training_output', # 模型训练输出路径 - evaluation_strategy="epoch", # 每个epoch结束后进行评估 - learning_rate=5e-5, # 学习率 - num_train_epochs=5, # 训练轮数 - warmup_steps=2000, # warmup步数 - per_device_train_batch_size=16, # 每个GPU的训练batch size - per_device_eval_batch_size=64, # 每个GPU的评估batch size - seed=120, # 随机数种子 - optim="adamw_torch", # 优化器 - save_strategy="epoch", # 保存模型的策略 - load_best_model_at_end=True, # 训练结束后加载最佳模型 - metric_for_best_model="accuracy", # 用于选择最佳模型的指标 + output_dir='./BERT_training_output', + evaluation_strategy="epoch", + learning_rate=5e-5, + num_train_epochs=8, + warmup_steps=2000, + per_device_train_batch_size=8, + per_device_eval_batch_size=32, + seed=120, + optim="adamw_torch", + save_strategy="epoch", + load_best_model_at_end=True, + metric_for_best_model='accuracy', ) + def preprocess_logits_for_metrics(logits, labels): if isinstance(logits, tuple): # Depending on the model and config, logits may contain extra tensors, @@ -68,34 +67,93 @@ def preprocess_logits_for_metrics(logits, labels): logits = logits[0] return logits +# 定义文件名 +output_dir = "./metrics/" +os.makedirs(output_dir, exist_ok=True) # 创建目录 +output_file = os.path.join(output_dir, "val_metrics.txt") + +file_labels = { + 'describe': 0, + 'expected': 1, + 'reproduce': 2, + 'actual': 3, + 'environment': 4, + 'additional': 5 +} + def compute_metrics(eval_pred): predictions, labels = eval_pred y_pred = np.argmax(predictions, axis=1) y_true = labels - f1_average = f1_score(y_true=y_true, y_pred=y_pred, average='binary') + + # 计算每个类别的指标 + f1_scores = f1_score(y_true, y_pred, average=None) accuracy = accuracy_score(y_true, y_pred) - precision = precision_score(y_true, y_pred, average='binary') - recall = recall_score(y_true, y_pred, average='binary') - # return as dictionary - metrics = { - 'f1': f1_average, - # 'ap': ap, - 'accuracy': accuracy, - 'precision': precision, - 'recall': recall + precision_scores = precision_score(y_true, y_pred, average=None) + recall_scores = recall_score(y_true, y_pred, average=None) + + # 计算总体指标 + overall_f1 = f1_score(y_true, y_pred, average='weighted') + overall_accuracy = accuracy_score(y_true, y_pred) + overall_precision = precision_score(y_true, y_pred, average='weighted') + overall_recall = recall_score(y_true, y_pred, average='weighted') + + # 将数字标签转换为对应的文字标签 + label_texts = {v: k for k, v in file_labels.items()} + + # 写入文件 + with open(output_file, "a") as f: + f.write("Class Metrics:\n") + for i in range(len(f1_scores)): + label_name = label_texts[i] # 使用文字标签 + f.write(f"Class: {label_name}\n") + f.write(f"F1 Score: {f1_scores[i]}\n") + f.write(f"Accuracy: {accuracy_score(y_true[y_true == i], y_pred[y_true == i])}\n") + f.write(f"Precision: {precision_scores[i]}\n") + f.write(f"Recall: {recall_scores[i]}\n") + f.write("\n") + + # 写入总体指标 + f.write("Overall Metrics:\n") + f.write(f"Overall F1 Score: {overall_f1}\n") + f.write(f"Overall Accuracy: {overall_accuracy}\n") + f.write(f"Overall Precision: {overall_precision}\n") + f.write(f"Overall Recall: {overall_recall}\n") + f.write("--------------------------------------------------\n") + + + # 保存每个类别的指标 + metrics = {} + for i in range(len(f1_scores)): + label_name = label_texts[i] # 使用文字标签 + metrics[label_name] = { + 'f1_score': f1_scores[i], + 'accuracy': accuracy_score(y_true[y_true == i], y_pred[y_true == i]), + 'precision': precision_scores[i], + 'recall': recall_scores[i] + } + + # 保存总体指标 + metrics['overall'] = { + 'f1_score': overall_f1, + 'accuracy': overall_accuracy, + 'precision': overall_precision, + 'recall': overall_recall } - return metrics + return { + 'accuracy': overall_accuracy, # 返回准确度作为评估指标 + } -# 定义Trainer类，用于训练模型 +# 创建Trainer实例 trainer = Trainer( - model=model, # 训练的模型 - args=training_args, # 训练参数 - train_dataset=train_dataset, # 训练集 - eval_dataset=val_dataset, # 验证集 - compute_metrics=compute_metrics, # 计算评估指标 + model=model, + args=training_args, + train_dataset=train_dataset, + eval_dataset=val_dataset, + compute_metrics=compute_metrics, preprocess_logits_for_metrics=preprocess_logits_for_metrics, # 预处理评估指标 ) -# 开始训练 -trainer.train() \ No newline at end of file +# 训练模型 +trainer.train()