Update bert_train.py
This commit is contained in:
parent
670a8dce17
commit
9cff4423ed
|
|
@ -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()
|
||||
# 训练模型
|
||||
trainer.train()
|
||||
|
|
|
|||
Loading…
Reference in New Issue