From ce0c6da36a0f447e3e7f316cba8e1c17c02ae8af Mon Sep 17 00:00:00 2001 From: ywd09 Date: Thu, 14 May 2026 00:37:04 +0800 Subject: [PATCH] =?UTF-8?q?=E6=A8=A1=E5=9E=8B=E8=AF=84=E4=BC=B0=E7=A8=8B?= =?UTF-8?q?=E5=BA=8F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- Curve.py | 50 ++++++++++++++++++ Evaluate.py | 147 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 197 insertions(+) create mode 100644 Curve.py create mode 100644 Evaluate.py diff --git a/Curve.py b/Curve.py new file mode 100644 index 0000000..7c38b37 --- /dev/null +++ b/Curve.py @@ -0,0 +1,50 @@ +""" +plot_training_curves.py +从 training_log.csv 读取日志，绘制 Loss / F1 / Accuracy / LR 曲线 +""" + +import pandas as pd +import matplotlib.pyplot as plt +import matplotlib + +matplotlib.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans'] +matplotlib.rcParams['axes.unicode_minus'] = False + +# ============ 读取数据 ============ +df = pd.read_csv('training_log.csv') +best_rows = df[df['best'] == 'best'] + +fig, axes = plt.subplots(2, 2, figsize=(14, 10)) + +# ---- 1. Loss ---- +ax = axes[0, 0] +ax.plot(df['epoch'], df['train_loss'], label='Train Loss', color='#1f77b4', lw=1.5) +ax.plot(df['epoch'], df['val_loss'], label='Val Loss', color='#ff7f0e', lw=1.5) +ax.set_xlabel('Epoch'); ax.set_ylabel('Loss'); ax.set_title('Loss vs Epoch') +ax.legend(); ax.grid(True, alpha=0.3) + +# ---- 2. F1 Score ---- +ax = axes[0, 1] +ax.plot(df['epoch'], df['train_f1'], label='Train F1', color='#1f77b4', lw=1.5) +ax.plot(df['epoch'], df['val_f1'], label='Val F1', color='#ff7f0e', lw=1.5) +ax.set_xlabel('Epoch'); ax.set_ylabel('F1 Score'); ax.set_title('F1 Score vs Epoch') +ax.legend(); ax.grid(True, alpha=0.3) + +# ---- 3. Accuracy ---- +ax = axes[1, 0] +ax.plot(df['epoch'], df['train_acc'], label='Train Acc', color='#1f77b4', lw=1.5) +ax.plot(df['epoch'], df['val_acc'], label='Val Acc', color='#ff7f0e', lw=1.5) +ax.set_xlabel('Epoch'); ax.set_ylabel('Accuracy (%)'); ax.set_title('Accuracy vs Epoch') +ax.legend(); ax.grid(True, alpha=0.3) + +# ---- 4. Learning Rate ---- +ax = axes[1, 1] +ax.plot(df['epoch'], df['lr'], color='#2ca02c', lw=1.5) +ax.set_xlabel('Epoch'); ax.set_ylabel('Learning Rate'); ax.set_title('Learning Rate vs Epoch') +ax.ticklabel_format(style='scientific', axis='y', scilimits=(0, 0)) +ax.grid(True, alpha=0.3) + +plt.tight_layout() +plt.savefig('training_curves.png', dpi=150, bbox_inches='tight') +plt.show() +print("训练曲线已保存: training_curves.png") diff --git a/Evaluate.py b/Evaluate.py new file mode 100644 index 0000000..7c78f1c --- /dev/null +++ b/Evaluate.py @@ -0,0 +1,147 @@ +""" +evaluate_and_plot.py +加载模型，在验证集上推理，绘制混淆矩阵 / ROC / PR 曲线 +""" + +import os +import numpy as np +import matplotlib.pyplot as plt +import matplotlib + +import torch +from torch.utils.data import DataLoader +from torchvision import transforms +from sklearn.metrics import ( + confusion_matrix, ConfusionMatrixDisplay, + roc_curve, auc, + precision_recall_curve, average_precision_score, +) + +from Model import Net +from Dataloader import RobustImageFolder + +matplotlib.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans'] +matplotlib.rcParams['axes.unicode_minus'] = False + +# ============================================================ +# ★★★ 需要你修改的参数 ★★★ +# ============================================================ +MODEL_PATH = 'best_model.pth' # 模型权重路径 +DATA_ROOT = '../trash_division_data/ultimate_4_class/' # 数据集根目录 +BATCH_SIZE = 32 +IMAGE_SIZE = 256 +NUM_WORKERS = 4 +# ============================================================ + +# ---------- 1. 加载验证集 ---------- +val_transform = transforms.Compose([ + transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)), + transforms.ToTensor(), + transforms.Normalize(mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225]), +]) + +val_dataset = RobustImageFolder( + root=os.path.join(DATA_ROOT, 'val'), + transform=val_transform, +) +val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, + shuffle=False, num_workers=NUM_WORKERS, + pin_memory=True, drop_last=False) + +class_names = val_dataset.classes +num_classes = len(class_names) +print(f"类别: {class_names}") + +# ---------- 2. 加载模型 ---------- +device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') +model = Net(num_classes=num_classes) +state_dict = torch.load(MODEL_PATH, map_location=device) +if 'model_state_dict' in state_dict: + state_dict = state_dict['model_state_dict'] +elif 'model' in state_dict: + state_dict = state_dict['model'] +model.load_state_dict(state_dict) +model = model.to(device).eval() +print("模型加载完成") + +# ---------- 3. 推理 ---------- +all_labels = [] +all_probs = [] + +with torch.no_grad(): + for images, labels in val_loader: + images = images.to(device) + probs = torch.softmax(model(images), dim=1) + all_labels.append(labels.numpy()) + all_probs.append(probs.cpu().numpy()) + +all_labels = np.concatenate(all_labels) +all_probs = np.concatenate(all_probs) +all_preds = np.argmax(all_probs, axis=1) +print(f"推理完成, 共 {len(all_labels)} 样本") + +# ============================================================ +# ① 混淆矩阵 +# ============================================================ +cm = confusion_matrix(all_labels, all_preds) +fig, ax = plt.subplots(figsize=(8, 7)) +ConfusionMatrixDisplay(cm, display_labels=class_names).plot( + ax=ax, cmap='Blues', values_format='d', xticks_rotation=30) +ax.set_title('Confusion Matrix', fontsize=14) +plt.tight_layout() +plt.savefig('confusion_matrix.png', dpi=150, bbox_inches='tight') +plt.show() +print("混淆矩阵已保存: confusion_matrix.png") + +# ============================================================ +# ② ROC 曲线 (One-vs-Rest + Macro-average) +# ============================================================ +one_hot = np.eye(num_classes)[all_labels] +colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728'] + +fig, ax = plt.subplots(figsize=(8, 7)) +fpr_d, tpr_d, auc_d = {}, {}, {} + +for i in range(num_classes): + fpr_d[i], tpr_d[i], _ = roc_curve(one_hot[:, i], all_probs[:, i]) + auc_d[i] = auc(fpr_d[i], tpr_d[i]) + ax.plot(fpr_d[i], tpr_d[i], color=colors[i], lw=2, + label=f'{class_names[i]} (AUC={auc_d[i]:.4f})') + +# Macro-average +all_fpr = np.unique(np.concatenate([fpr_d[i] for i in range(num_classes)])) +mean_tpr = sum(np.interp(all_fpr, fpr_d[i], tpr_d[i]) for i in range(num_classes)) / num_classes +macro_auc = auc(all_fpr, mean_tpr) +ax.plot(all_fpr, mean_tpr, 'navy', lw=2, ls='--', + label=f'Macro-avg (AUC={macro_auc:.4f})') +ax.plot([0, 1], [0, 1], 'k--', lw=1, alpha=0.5) + +ax.set_xlim(0, 1); ax.set_ylim(0, 1.05) +ax.set_xlabel('False Positive Rate'); ax.set_ylabel('True Positive Rate') +ax.set_title('ROC Curve', fontsize=14) +ax.legend(loc='lower right'); ax.grid(True, alpha=0.3) +plt.tight_layout() +plt.savefig('roc_curve.png', dpi=150, bbox_inches='tight') +plt.show() +print("ROC 曲线已保存: roc_curve.png") + +# ============================================================ +# ③ Precision-Recall 曲线 +# ============================================================ +fig, ax = plt.subplots(figsize=(8, 7)) + +for i in range(num_classes): + prec, rec, _ = precision_recall_curve(one_hot[:, i], all_probs[:, i]) + ap = average_precision_score(one_hot[:, i], all_probs[:, i]) + ax.plot(rec, prec, color=colors[i], lw=2, + label=f'{class_names[i]} (AP={ap:.4f})') + +ax.set_xlim(0, 1); ax.set_ylim(0, 1.05) +ax.set_xlabel('Recall'); ax.set_ylabel('Precision') +ax.set_title('Precision-Recall Curve', fontsize=14) +ax.legend(loc='best'); ax.grid(True, alpha=0.3) +plt.tight_layout() +plt.savefig('pr_curve.png', dpi=150, bbox_inches='tight') +plt.show() +print("PR 曲线已保存: pr_curve.png")