improve training for class-imbalanced DR classification

- WeightedRandomSampler for balanced class sampling during training
- FocalLoss (gamma=2) for hard-example mining, no alpha (sampler handles balance)
- Stronger augmentation: RandomResizedCrop, ColorJitter, GaussianBlur, 30deg rotation
- Discriminative LR: backbone at 0.1x, head at 1x
- 3-epoch linear warmup + cosine decay scheduler
- Early stopping on macro F1 instead of weighted F1
- Updated defaults: batch_size=64, epochs=50, patience=10
- Comparison chart now shows both macro and weighted F1
- DataParallel checkpoint handling (strip module. prefix)
- Enable cudnn.benchmark for training speed

1 files changed, 22 insertions, 12 deletions

diff --git a/evaluate.py b/evaluate.py
index 7a39db4..7f996dc 100644
--- a/evaluate.py
+++ b/evaluate.py
@@ -76,10 +76,10 @@ def plot_training_curves(
     ax1.set_title('Training and Validation Loss')
     ax1.legend()
 
-    ax2.plot(epochs, val_f1s, label='Val weighted F1', color='green')
+    ax2.plot(epochs, val_f1s, label='Val macro F1', color='green')
     ax2.set_xlabel('Epoch')
-    ax2.set_ylabel('Weighted F1')
-    ax2.set_title('Validation Weighted F1')
+    ax2.set_ylabel('Macro F1')
+    ax2.set_title('Validation Macro F1')
     ax2.legend()
 
     plt.tight_layout()
@@ -127,18 +127,28 @@ def generate_summary(all_results: dict, save_dir: str) -> None:
         writer.writerows(rows)
     print(f"\nSummary saved to {csv_path}")
 
-    # Comparison bar chart
-    labels = [f"{r['model']}\nCLAHE={r['clahe']}" for r in rows]
+    # Comparison bar chart — show both macro and weighted F1
+    x_labels = [f"{r['model']}\nCLAHE={r['clahe']}" for r in rows]
     w_f1_vals = [r['weighted_f1'] for r in rows]
-
-    fig, ax = plt.subplots(figsize=(10, 5))
-    bars = ax.bar(labels, w_f1_vals, color=['#4C72B0', '#DD8452'] * 3)
-    ax.set_ylabel('Weighted F1-Score')
-    ax.set_title('Weighted F1 Comparison Across Experiments')
+    m_f1_vals = [r['macro_f1'] for r in rows]
+
+    x = np.arange(len(x_labels))
+    bar_width = 0.35
+    fig, ax = plt.subplots(figsize=(12, 5))
+    bars_w = ax.bar(x - bar_width / 2, w_f1_vals, bar_width, label='Weighted F1', color='#4C72B0')
+    bars_m = ax.bar(x + bar_width / 2, m_f1_vals, bar_width, label='Macro F1', color='#DD8452')
+    ax.set_ylabel('F1-Score')
+    ax.set_title('F1 Comparison Across Experiments')
+    ax.set_xticks(x)
+    ax.set_xticklabels(x_labels)
     ax.set_ylim(0, 1.0)
-    for bar, val in zip(bars, w_f1_vals):
+    ax.legend()
+    for bar, val in zip(bars_w, w_f1_vals):
+        ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.01,
+                f'{val:.3f}', ha='center', va='bottom', fontsize=8)
+    for bar, val in zip(bars_m, m_f1_vals):
         ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.01,
-                f'{val:.4f}', ha='center', va='bottom', fontsize=9)
+                f'{val:.3f}', ha='center', va='bottom', fontsize=8)
     plt.tight_layout()
     chart_path = os.path.join(save_dir, 'comparison_chart.png')
     fig.savefig(chart_path, dpi=150)