Upload 3 files
Browse files
Appendix - Export Pytorch to ONNX.ipynb
ADDED
The diff for this file is too large to render.
See raw diff
|
|
OModeling-2.py
ADDED
@@ -0,0 +1,1894 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import gc
|
3 |
+
import json
|
4 |
+
import math
|
5 |
+
import torch
|
6 |
+
import mlflow
|
7 |
+
import logging
|
8 |
+
import platform
|
9 |
+
import numpy as np
|
10 |
+
import pandas as pd
|
11 |
+
from PIL import Image
|
12 |
+
from tqdm import tqdm
|
13 |
+
import torch.nn as nn
|
14 |
+
import torch.optim as optim
|
15 |
+
from torchvision import models
|
16 |
+
import matplotlib.pyplot as plt
|
17 |
+
import torch.nn.functional as F
|
18 |
+
from sklearn.manifold import TSNE
|
19 |
+
from torchvision import transforms
|
20 |
+
from kymatio.torch import Scattering2D
|
21 |
+
from torch.utils.data import Dataset, DataLoader
|
22 |
+
from pytorch_metric_learning.miners import BatchHardMiner
|
23 |
+
from pytorch_metric_learning.losses import MultiSimilarityLoss
|
24 |
+
from torch.optim.lr_scheduler import CosineAnnealingLR, ReduceLROnPlateau
|
25 |
+
from sklearn.metrics import roc_curve, auc, precision_recall_fscore_support
|
26 |
+
from typing import Dict, List, Tuple, Optional, Union, Any
|
27 |
+
from dataclasses import dataclass, asdict
|
28 |
+
import warnings
|
29 |
+
warnings.filterwarnings('ignore')
|
30 |
+
|
31 |
+
# ----------------------------
|
32 |
+
# Configuration Management
|
33 |
+
# ----------------------------
|
34 |
+
@dataclass
|
35 |
+
@dataclass
|
36 |
+
class TrainingConfig:
|
37 |
+
|
38 |
+
# Model Architecture
|
39 |
+
model_name: str = "resnet34"
|
40 |
+
embedding_dim: int = 128
|
41 |
+
normalize_embeddings: bool = True
|
42 |
+
pretrained_path: Optional[str] = "../../model/pretrained_model/ResNet34.pt"
|
43 |
+
|
44 |
+
# Training Hyperparameters
|
45 |
+
batch_size: int = 512
|
46 |
+
max_epochs: int = 20
|
47 |
+
grad_accum_steps: int = 10
|
48 |
+
device: str = "cuda" if torch.cuda.is_available() else "cpu"
|
49 |
+
|
50 |
+
# Learning Rate Configuration
|
51 |
+
head_lr: float = 1e-3 # Higher LR for embedding head (untrained)
|
52 |
+
backbone_lr: float = 1e-4 # Lower LR for backbone (pretrained)
|
53 |
+
lr_scheduler: str = "cosine" # "cosine" or "plateau"
|
54 |
+
weight_decay: float = 1e-4
|
55 |
+
|
56 |
+
# Curriculum Learning Parameters - ADJUSTED FOR PRECISION
|
57 |
+
curriculum_strategy: str = "progressive" # "progressive", "exponential", "linear"
|
58 |
+
initial_hard_ratio: float = 0.6 # Increased from 0.1 for more hard negatives early
|
59 |
+
final_hard_ratio: float = 0.9 # Increased from 0.8 for focus on hard cases
|
60 |
+
curriculum_warmup_epochs: int = 1 # Reduced from 2 for faster hard sample exposure
|
61 |
+
|
62 |
+
# Data Augmentation
|
63 |
+
remove_bg: bool = False
|
64 |
+
augmentation_strength: float = 0.5 # 0.0 = no aug, 1.0 = strong aug
|
65 |
+
|
66 |
+
# Loss Configuration - ADJUSTED FOR PRECISION
|
67 |
+
multisim_alpha: float = 2.5 # Increased from 2.0 (penalize false positives more)
|
68 |
+
multisim_beta: float = 60.0 # Increased from 50.0 (larger margin)
|
69 |
+
multisim_base: float = 0.4 # Decreased from 0.5 (stricter similarity)
|
70 |
+
|
71 |
+
# Triplet Loss Parameters - NEW
|
72 |
+
triplet_margin: float = 1.0 # Margin for triplet loss
|
73 |
+
triplet_weight: float = 0.3 # Weight for triplet loss component
|
74 |
+
false_positive_penalty_weight: float = 0.3 # Extra penalty for false positives
|
75 |
+
|
76 |
+
# Mining Configuration
|
77 |
+
use_hard_mining: bool = True
|
78 |
+
|
79 |
+
# Precision Focus Parameters - NEW
|
80 |
+
target_precision: float = 0.75 # Target precision for threshold selection
|
81 |
+
negative_weight_multiplier: float = 2.5 # How much more to weight hard negatives
|
82 |
+
|
83 |
+
# Checkpoint Configuration
|
84 |
+
run_id: Optional[str] = None
|
85 |
+
last_epoch_weights: Optional[str] = None
|
86 |
+
save_frequency: int = 1 # Save every N epochs
|
87 |
+
|
88 |
+
# Early Stopping
|
89 |
+
patience: int = 15
|
90 |
+
min_delta: float = 0.001
|
91 |
+
|
92 |
+
# Logging
|
93 |
+
log_frequency: int = 100 # Log every N steps
|
94 |
+
visualize_frequency: int = 1 # Visualize every N epochs
|
95 |
+
|
96 |
+
tracking_uri: str = "http://127.0.0.1:5555"
|
97 |
+
|
98 |
+
def __post_init__(self):
|
99 |
+
"""Validate configuration parameters."""
|
100 |
+
assert 0.0 <= self.initial_hard_ratio <= 1.0, "Initial hard ratio must be in [0, 1]"
|
101 |
+
assert 0.0 <= self.final_hard_ratio <= 1.0, "Final hard ratio must be in [0, 1]"
|
102 |
+
assert self.curriculum_strategy in ["progressive", "exponential", "linear"]
|
103 |
+
assert self.lr_scheduler in ["cosine", "plateau"]
|
104 |
+
assert 0.0 <= self.target_precision <= 1.0, "Target precision must be in [0, 1]"
|
105 |
+
|
106 |
+
|
107 |
+
# Global configuration
|
108 |
+
CONFIG = TrainingConfig()
|
109 |
+
|
110 |
+
# ----------------------------
|
111 |
+
# MLFlow Setup
|
112 |
+
# ----------------------------
|
113 |
+
class MLFlowManager:
|
114 |
+
"""Centralized MLflow management for experiment tracking."""
|
115 |
+
|
116 |
+
def __init__(self, tracking_uri: str = "http://127.0.0.1:5555"):
|
117 |
+
mlflow.set_tracking_uri(tracking_uri)
|
118 |
+
self.experiment_name = "Signature Verification - Advanced Training"
|
119 |
+
self._setup_experiment()
|
120 |
+
|
121 |
+
def _setup_experiment(self):
|
122 |
+
"""Setup MLflow experiment."""
|
123 |
+
try:
|
124 |
+
self.experiment_id = mlflow.create_experiment(self.experiment_name)
|
125 |
+
except:
|
126 |
+
self.experiment_id = mlflow.get_experiment_by_name(self.experiment_name).experiment_id
|
127 |
+
|
128 |
+
def start_run(self, run_id: Optional[str] = None):
|
129 |
+
"""Start MLflow run with configuration logging."""
|
130 |
+
return mlflow.start_run(run_id=run_id, experiment_id=self.experiment_id)
|
131 |
+
|
132 |
+
def log_config(self, config: TrainingConfig):
|
133 |
+
"""Log training configuration."""
|
134 |
+
config_dict = asdict(config)
|
135 |
+
mlflow.log_params(config_dict)
|
136 |
+
|
137 |
+
# ----------------------------
|
138 |
+
# Curriculum Learning Manager
|
139 |
+
# ----------------------------
|
140 |
+
class CurriculumLearningManager:
|
141 |
+
"""Advanced curriculum learning for both hard positives and hard negatives."""
|
142 |
+
|
143 |
+
def __init__(self, config: TrainingConfig):
|
144 |
+
self.config = config
|
145 |
+
self.current_epoch = 0
|
146 |
+
|
147 |
+
def get_hard_ratio(self, epoch: int) -> float:
|
148 |
+
"""Get hard negative ratio (forgeries) for current epoch."""
|
149 |
+
if epoch < self.config.curriculum_warmup_epochs:
|
150 |
+
return self.config.initial_hard_ratio
|
151 |
+
|
152 |
+
# Target: reach final_hard_ratio by max_epochs // 2
|
153 |
+
target_epoch = max(self.config.max_epochs // 2, self.config.curriculum_warmup_epochs + 3)
|
154 |
+
|
155 |
+
if epoch >= target_epoch:
|
156 |
+
return self.config.final_hard_ratio
|
157 |
+
|
158 |
+
# Aggressive progression to reach target by mid-training
|
159 |
+
progress = (epoch - self.config.curriculum_warmup_epochs) / (target_epoch - self.config.curriculum_warmup_epochs)
|
160 |
+
|
161 |
+
initial = self.config.initial_hard_ratio
|
162 |
+
final = self.config.final_hard_ratio
|
163 |
+
|
164 |
+
if self.config.curriculum_strategy == "progressive":
|
165 |
+
# Very aggressive: exponential growth early, then plateau
|
166 |
+
ratio = initial + (final - initial) * (progress ** 0.5)
|
167 |
+
elif self.config.curriculum_strategy == "exponential":
|
168 |
+
ratio = initial + (final - initial) * (progress ** 0.3)
|
169 |
+
else: # linear
|
170 |
+
ratio = initial + (final - initial) * progress
|
171 |
+
|
172 |
+
return min(max(ratio, 0.0), 1.0)
|
173 |
+
|
174 |
+
def get_hard_positive_ratio(self, epoch: int) -> float:
|
175 |
+
"""Get hard positive ratio for current epoch - increases more gradually."""
|
176 |
+
if epoch < self.config.curriculum_warmup_epochs:
|
177 |
+
return 0.1 # Start with 10% hard positives
|
178 |
+
|
179 |
+
# Hard positives should increase more gradually than hard negatives
|
180 |
+
max_epochs = self.config.max_epochs
|
181 |
+
progress = min(1.0, (epoch - self.config.curriculum_warmup_epochs) / (max_epochs - self.config.curriculum_warmup_epochs))
|
182 |
+
|
183 |
+
# Target 40% hard positives by end of training
|
184 |
+
initial_ratio = 0.1
|
185 |
+
final_ratio = 0.4
|
186 |
+
|
187 |
+
if self.config.curriculum_strategy == "progressive":
|
188 |
+
ratio = initial_ratio + (final_ratio - initial_ratio) * (progress ** 0.7)
|
189 |
+
else:
|
190 |
+
ratio = initial_ratio + (final_ratio - initial_ratio) * progress
|
191 |
+
|
192 |
+
return min(max(ratio, 0.0), final_ratio)
|
193 |
+
|
194 |
+
def get_mining_difficulty(self, epoch: int) -> Dict[str, float]:
|
195 |
+
"""Adaptive mining parameters for both hard positives and negatives."""
|
196 |
+
progress = min(1.0, epoch / self.config.max_epochs)
|
197 |
+
|
198 |
+
# Separate ratios for hard positives and hard negatives
|
199 |
+
hard_negative_ratio = self.get_hard_ratio(epoch)
|
200 |
+
hard_positive_ratio = self.get_hard_positive_ratio(epoch)
|
201 |
+
|
202 |
+
# Dynamic weights for different sample types
|
203 |
+
hard_pos_weight = 1.0 + 2.0 * progress # 1.0 → 3.0
|
204 |
+
hard_neg_weight = 1.0 + 4.0 * progress # 1.0 → 5.0 (harder negatives more important)
|
205 |
+
|
206 |
+
return {
|
207 |
+
# Margin parameters
|
208 |
+
"margin_multiplier": 1.0 + 0.5 * progress,
|
209 |
+
|
210 |
+
# Hard sample ratios
|
211 |
+
"hard_negative_ratio": hard_negative_ratio,
|
212 |
+
"hard_positive_ratio": hard_positive_ratio,
|
213 |
+
"current_hard_ratio": hard_negative_ratio, # For backward compatibility
|
214 |
+
|
215 |
+
# Sample weights
|
216 |
+
"hard_positive_weight": hard_pos_weight,
|
217 |
+
"hard_negative_weight": hard_neg_weight,
|
218 |
+
"semi_positive_weight": 1.0 + 1.0 * progress,
|
219 |
+
"semi_negative_weight": 1.0 + 2.0 * progress,
|
220 |
+
|
221 |
+
# Difficulty thresholds
|
222 |
+
"difficulty_threshold": 0.05 + 0.15 * progress,
|
223 |
+
"selectivity": 0.8 + 0.2 * progress,
|
224 |
+
|
225 |
+
# Mining aggressiveness
|
226 |
+
"mining_temperature": max(0.5, 1.0 - 0.5 * progress), # Decreases over time
|
227 |
+
|
228 |
+
# Focus balance (0 = equal focus, 1 = focus on negatives)
|
229 |
+
"negative_focus": 0.5 + 0.3 * progress
|
230 |
+
}
|
231 |
+
# ----------------------------
|
232 |
+
# Enhanced Dataset with Advanced Curriculum Learning
|
233 |
+
# ----------------------------
|
234 |
+
class SignatureDataset(Dataset):
|
235 |
+
"""
|
236 |
+
Advanced signature dataset with curriculum learning and mining statistics.
|
237 |
+
"""
|
238 |
+
|
239 |
+
def __init__(
|
240 |
+
self,
|
241 |
+
folder_img: str,
|
242 |
+
excel_data: pd.DataFrame,
|
243 |
+
curriculum_manager: CurriculumLearningManager,
|
244 |
+
transform: Optional[transforms.Compose] = None,
|
245 |
+
is_train: bool = True,
|
246 |
+
config: TrainingConfig = CONFIG
|
247 |
+
):
|
248 |
+
self.folder_img = folder_img
|
249 |
+
self.is_train = is_train
|
250 |
+
self.config = config
|
251 |
+
self.curriculum_manager = curriculum_manager
|
252 |
+
self.transform = transform or self._default_transforms()
|
253 |
+
self.excel_data = excel_data.reset_index(drop=True)
|
254 |
+
self.current_epoch = 0
|
255 |
+
|
256 |
+
# Data preparation
|
257 |
+
self._handle_excel_person_ids()
|
258 |
+
self._categorize_difficulty()
|
259 |
+
|
260 |
+
# Curriculum learning data
|
261 |
+
self.epoch_data = []
|
262 |
+
self._prepare_epoch_data()
|
263 |
+
def _handle_excel_person_ids(self):
|
264 |
+
"""Properly separate genuine vs forged signature IDs with compact offset."""
|
265 |
+
# Map genuine person IDs to 0, 1, 2, ...
|
266 |
+
genuine_ids = pd.concat([
|
267 |
+
self.excel_data["anchor_id"],
|
268 |
+
self.excel_data[self.excel_data["easy_or_hard"] == "easy"]["negative_id"]
|
269 |
+
]).unique()
|
270 |
+
|
271 |
+
self.genuine_id_mapping = {val: idx for idx, val in enumerate(genuine_ids)}
|
272 |
+
max_genuine_id = len(genuine_ids)
|
273 |
+
|
274 |
+
# Create forgery ID space with SMALLER offset (just enough to avoid collisions)
|
275 |
+
forged_data = self.excel_data[self.excel_data["easy_or_hard"] == "hard"]
|
276 |
+
if len(forged_data) > 0:
|
277 |
+
unique_forged_persons = forged_data["negative_id"].unique()
|
278 |
+
self.forgery_id_mapping = {
|
279 |
+
val: idx + max_genuine_id + 100 # Smaller offset: 100 instead of 1000
|
280 |
+
for idx, val in enumerate(unique_forged_persons)
|
281 |
+
}
|
282 |
+
else:
|
283 |
+
self.forgery_id_mapping = {}
|
284 |
+
|
285 |
+
# Apply mappings
|
286 |
+
self.excel_data["anchor_id"] = self.excel_data["anchor_id"].map(self.genuine_id_mapping)
|
287 |
+
|
288 |
+
# Handle negatives based on type
|
289 |
+
new_negative_ids = []
|
290 |
+
for idx, row in self.excel_data.iterrows():
|
291 |
+
if row["easy_or_hard"] == "easy":
|
292 |
+
# Genuine different person: use regular ID
|
293 |
+
new_negative_ids.append(self.genuine_id_mapping[row["negative_id"]])
|
294 |
+
else:
|
295 |
+
# Forged signature: use offset ID to prevent clustering with genuine
|
296 |
+
new_negative_ids.append(self.forgery_id_mapping[row["negative_id"]])
|
297 |
+
|
298 |
+
self.excel_data["negative_id"] = new_negative_ids
|
299 |
+
|
300 |
+
print(f"ID mapping: Genuine IDs 0-{max_genuine_id-1}, Forgery IDs {max_genuine_id+100}-{max_genuine_id+100+len(self.forgery_id_mapping)-1}")
|
301 |
+
|
302 |
+
|
303 |
+
def _categorize_difficulty(self):
|
304 |
+
"""Categorize samples by difficulty if not already done."""
|
305 |
+
if self.is_train and "easy_or_hard" in self.excel_data.columns:
|
306 |
+
self.easy_df = self.excel_data[self.excel_data["easy_or_hard"] == "easy"]
|
307 |
+
self.hard_df = self.excel_data[self.excel_data["easy_or_hard"] == "hard"]
|
308 |
+
else:
|
309 |
+
# All samples treated as medium difficulty
|
310 |
+
self.easy_df = self.excel_data
|
311 |
+
self.hard_df = pd.DataFrame() # Empty hard samples
|
312 |
+
|
313 |
+
def _prepare_epoch_data(self):
|
314 |
+
"""Prepare data for current epoch based on curriculum."""
|
315 |
+
if not self.is_train:
|
316 |
+
# Validation data preparation with better error handling
|
317 |
+
if "image_1_path" in self.excel_data.columns and "image_2_path" in self.excel_data.columns:
|
318 |
+
# Standard pair format
|
319 |
+
required_cols = ["image_1_path", "image_2_path", "label"]
|
320 |
+
|
321 |
+
# Find ID columns
|
322 |
+
id_cols = [col for col in self.excel_data.columns if "id" in col.lower()]
|
323 |
+
if len(id_cols) >= 2:
|
324 |
+
required_cols.extend(id_cols[-2:]) # Take last 2 ID columns
|
325 |
+
else:
|
326 |
+
# Create dummy IDs if none exist
|
327 |
+
self.excel_data["dummy_id1"] = 0
|
328 |
+
self.excel_data["dummy_id2"] = 1
|
329 |
+
required_cols.extend(["dummy_id1", "dummy_id2"])
|
330 |
+
|
331 |
+
self.epoch_data = self.excel_data[required_cols].values.tolist()
|
332 |
+
|
333 |
+
else:
|
334 |
+
# Fallback: try to use all available columns
|
335 |
+
print(f"Warning: Expected validation columns not found. Available: {list(self.excel_data.columns)}")
|
336 |
+
self.epoch_data = self.excel_data.values.tolist()
|
337 |
+
|
338 |
+
print(f"Validation data prepared: {len(self.epoch_data)} samples")
|
339 |
+
return
|
340 |
+
|
341 |
+
# Training data preparation (unchanged)
|
342 |
+
hard_ratio = self.curriculum_manager.get_hard_ratio(self.current_epoch)
|
343 |
+
|
344 |
+
if len(self.hard_df) > 0:
|
345 |
+
n_total = len(self.excel_data)
|
346 |
+
n_hard = int(n_total * hard_ratio)
|
347 |
+
n_easy = n_total - n_hard
|
348 |
+
|
349 |
+
hard_sample = self.hard_df.sample(
|
350 |
+
n=min(n_hard, len(self.hard_df)),
|
351 |
+
random_state=self.current_epoch,
|
352 |
+
replace=(n_hard > len(self.hard_df))
|
353 |
+
)
|
354 |
+
easy_sample = self.easy_df.sample(
|
355 |
+
n=min(n_easy, len(self.easy_df)),
|
356 |
+
random_state=self.current_epoch,
|
357 |
+
replace=(n_easy > len(self.easy_df))
|
358 |
+
)
|
359 |
+
|
360 |
+
epoch_df = pd.concat([hard_sample, easy_sample]).sample(
|
361 |
+
frac=1, random_state=self.current_epoch
|
362 |
+
).reset_index(drop=True)
|
363 |
+
|
364 |
+
print(f"Epoch {self.current_epoch}: {len(hard_sample)} hard + {len(easy_sample)} easy = {len(epoch_df)} total (target ratio: {hard_ratio:.2f})")
|
365 |
+
else:
|
366 |
+
epoch_df = self.excel_data.sample(
|
367 |
+
frac=1, random_state=self.current_epoch
|
368 |
+
).reset_index(drop=True)
|
369 |
+
|
370 |
+
required_cols = ["anchor_path", "positive_path", "negative_path", "anchor_id", "negative_id"]
|
371 |
+
missing_cols = [col for col in required_cols if col not in epoch_df.columns]
|
372 |
+
if missing_cols:
|
373 |
+
raise ValueError(f"Missing required training columns: {missing_cols}")
|
374 |
+
|
375 |
+
self.epoch_data = epoch_df[required_cols].values.tolist()
|
376 |
+
|
377 |
+
def set_epoch(self, epoch: int):
|
378 |
+
"""Update epoch and regenerate data."""
|
379 |
+
self.current_epoch = epoch
|
380 |
+
self._prepare_epoch_data()
|
381 |
+
|
382 |
+
def get_curriculum_stats(self) -> Dict[str, Any]:
|
383 |
+
"""Get current curriculum learning statistics."""
|
384 |
+
hard_ratio = self.curriculum_manager.get_hard_ratio(self.current_epoch)
|
385 |
+
mining_params = self.curriculum_manager.get_mining_difficulty(self.current_epoch)
|
386 |
+
|
387 |
+
return {
|
388 |
+
"epoch": self.current_epoch,
|
389 |
+
"hard_ratio": hard_ratio,
|
390 |
+
"easy_ratio": 1.0 - hard_ratio,
|
391 |
+
"total_samples": len(self.epoch_data),
|
392 |
+
**mining_params
|
393 |
+
}
|
394 |
+
|
395 |
+
def __len__(self) -> int:
|
396 |
+
return len(self.epoch_data)
|
397 |
+
|
398 |
+
def __getitem__(self, index: int) -> Tuple[torch.Tensor, ...]:
|
399 |
+
if self.is_train:
|
400 |
+
return self._get_train_item(index)
|
401 |
+
else:
|
402 |
+
return self._get_val_item(index)
|
403 |
+
|
404 |
+
def _get_train_item(self, index: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int]:
|
405 |
+
"""Return triplet: anchor, positive, negative with their IDs."""
|
406 |
+
anchor_path, positive_path, negative_path, pid, nid = self.epoch_data[index]
|
407 |
+
|
408 |
+
anchor = self._load_image(anchor_path)
|
409 |
+
positive = self._load_image(positive_path)
|
410 |
+
negative = self._load_image(negative_path)
|
411 |
+
|
412 |
+
return anchor, positive, negative, int(pid), int(nid)
|
413 |
+
|
414 |
+
def _get_val_item(self, index: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int]:
|
415 |
+
"""Return: img1, img2, label, id1, id2."""
|
416 |
+
data_row = self.epoch_data[index]
|
417 |
+
|
418 |
+
# Handle different data formats robustly
|
419 |
+
if len(data_row) >= 5:
|
420 |
+
img1_path, img2_path, label, id1, id2 = data_row[:5]
|
421 |
+
elif len(data_row) >= 3:
|
422 |
+
img1_path, img2_path, label = data_row[:3]
|
423 |
+
# Fallback IDs
|
424 |
+
id1, id2 = 0, 1
|
425 |
+
else:
|
426 |
+
raise ValueError(f"Invalid validation data format: expected at least 3 columns, got {len(data_row)}")
|
427 |
+
|
428 |
+
try:
|
429 |
+
img1 = self._load_image(img1_path)
|
430 |
+
img2 = self._load_image(img2_path)
|
431 |
+
|
432 |
+
return img1, img2, torch.tensor(float(label), dtype=torch.float32), int(id1), int(id2)
|
433 |
+
except Exception as e:
|
434 |
+
print(f"Error loading validation item {index}: {e}")
|
435 |
+
print(f"Data row: {data_row}")
|
436 |
+
raise
|
437 |
+
|
438 |
+
def _load_image(self, path: str) -> torch.Tensor:
|
439 |
+
"""Load and transform image."""
|
440 |
+
image = replace_background_with_white(
|
441 |
+
path, self.folder_img, remove_bg=self.config.remove_bg
|
442 |
+
)
|
443 |
+
return self.transform(image) if self.transform else image
|
444 |
+
|
445 |
+
def _default_transforms(self) -> transforms.Compose:
|
446 |
+
"""Get default transforms with configurable augmentation strength."""
|
447 |
+
normalize = transforms.Normalize(
|
448 |
+
mean=[0.485, 0.456, 0.406],
|
449 |
+
std=[0.229, 0.224, 0.225]
|
450 |
+
)
|
451 |
+
|
452 |
+
if self.is_train:
|
453 |
+
aug_strength = self.config.augmentation_strength
|
454 |
+
return transforms.Compose([
|
455 |
+
transforms.Resize((224, 224)),
|
456 |
+
transforms.RandomHorizontalFlip(p=0.5 * aug_strength),
|
457 |
+
transforms.RandomRotation(degrees=int(10 * aug_strength)),
|
458 |
+
transforms.ColorJitter(
|
459 |
+
brightness=0.2 * aug_strength,
|
460 |
+
contrast=0.2 * aug_strength
|
461 |
+
),
|
462 |
+
transforms.GaussianBlur(kernel_size=5, sigma=(0.1, 2.0 * aug_strength)),
|
463 |
+
transforms.ToTensor(),
|
464 |
+
normalize
|
465 |
+
])
|
466 |
+
|
467 |
+
return transforms.Compose([
|
468 |
+
transforms.Resize((224, 224)),
|
469 |
+
transforms.ToTensor(),
|
470 |
+
normalize
|
471 |
+
])
|
472 |
+
|
473 |
+
# ----------------------------
|
474 |
+
# Enhanced Model Architecture
|
475 |
+
# ----------------------------
|
476 |
+
class ResNetBackbone(nn.Module):
|
477 |
+
"""Enhanced ResNet backbone with better weight loading."""
|
478 |
+
|
479 |
+
def __init__(self, model_name: str = "resnet34", pretrained_path: Optional[str] = None):
|
480 |
+
super().__init__()
|
481 |
+
|
482 |
+
# Initialize the ResNet model
|
483 |
+
if model_name == "resnet18":
|
484 |
+
self.resnet = models.resnet18(weights=None)
|
485 |
+
elif model_name == "resnet34":
|
486 |
+
self.resnet = models.resnet34(weights=None)
|
487 |
+
elif model_name == "resnet50":
|
488 |
+
self.resnet = models.resnet50(weights=None)
|
489 |
+
else:
|
490 |
+
raise ValueError(f"Unsupported model_name: {model_name}")
|
491 |
+
|
492 |
+
# Load pretrained weights
|
493 |
+
if pretrained_path and os.path.exists(pretrained_path):
|
494 |
+
self._load_pretrained_weights(pretrained_path)
|
495 |
+
elif pretrained_path:
|
496 |
+
print(f"Warning: Pretrained path {pretrained_path} not found, using random initialization")
|
497 |
+
|
498 |
+
# Remove the fully connected layer
|
499 |
+
self.resnet.fc = nn.Identity()
|
500 |
+
|
501 |
+
# Get output dimension
|
502 |
+
with torch.no_grad():
|
503 |
+
dummy = torch.randn(1, 3, 224, 224)
|
504 |
+
self.output_dim = self.resnet(dummy).shape[1]
|
505 |
+
|
506 |
+
def _load_pretrained_weights(self, pretrained_path: str):
|
507 |
+
"""Load pretrained weights with comprehensive error handling."""
|
508 |
+
try:
|
509 |
+
checkpoint = torch.load(pretrained_path, map_location="cpu", weights_only=False)
|
510 |
+
state_dict = checkpoint.get("state_dict", checkpoint)
|
511 |
+
|
512 |
+
# Handle prefix issues
|
513 |
+
if not any(key.startswith("resnet.") for key in state_dict.keys()):
|
514 |
+
state_dict = {f"resnet.{k}": v for k, v in state_dict.items()}
|
515 |
+
|
516 |
+
# Filter matching keys and sizes
|
517 |
+
model_dict = self.state_dict()
|
518 |
+
filtered_state_dict = {
|
519 |
+
k: v for k, v in state_dict.items()
|
520 |
+
if k in model_dict and v.size() == model_dict[k].size()
|
521 |
+
}
|
522 |
+
|
523 |
+
# Load filtered weights
|
524 |
+
missing_keys = self.load_state_dict(filtered_state_dict, strict=False)
|
525 |
+
|
526 |
+
print(f"[INFO] Loaded pretrained weights: {len(filtered_state_dict)}/{len(model_dict)} parameters")
|
527 |
+
if missing_keys.missing_keys:
|
528 |
+
print(f"[INFO] Missing keys: {len(missing_keys.missing_keys)}")
|
529 |
+
|
530 |
+
except Exception as e:
|
531 |
+
print(f"[ERROR] Failed to load pretrained weights: {e}")
|
532 |
+
raise
|
533 |
+
|
534 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
535 |
+
return self.resnet(x)
|
536 |
+
|
537 |
+
class AdvancedEmbeddingHead(nn.Module):
|
538 |
+
"""Advanced embedding head with residual connections and normalization."""
|
539 |
+
|
540 |
+
def __init__(self, input_dim: int, embedding_dim: int, dropout: float = 0.5):
|
541 |
+
super().__init__()
|
542 |
+
|
543 |
+
self.input_dim = input_dim
|
544 |
+
self.embedding_dim = embedding_dim
|
545 |
+
|
546 |
+
# Multi-layer embedding head with residual connections
|
547 |
+
if input_dim > embedding_dim * 4:
|
548 |
+
hidden_dim = max(embedding_dim * 2, input_dim // 4)
|
549 |
+
self.layers = nn.Sequential(
|
550 |
+
nn.Linear(input_dim, hidden_dim),
|
551 |
+
nn.LayerNorm(hidden_dim),
|
552 |
+
nn.GELU(),
|
553 |
+
nn.Dropout(dropout),
|
554 |
+
|
555 |
+
nn.Linear(hidden_dim, embedding_dim * 2),
|
556 |
+
nn.LayerNorm(embedding_dim * 2),
|
557 |
+
nn.GELU(),
|
558 |
+
nn.Dropout(dropout / 2),
|
559 |
+
|
560 |
+
nn.Linear(embedding_dim * 2, embedding_dim),
|
561 |
+
nn.LayerNorm(embedding_dim)
|
562 |
+
)
|
563 |
+
else:
|
564 |
+
# Simple head for smaller dimensions
|
565 |
+
self.layers = nn.Sequential(
|
566 |
+
nn.Linear(input_dim, embedding_dim),
|
567 |
+
nn.LayerNorm(embedding_dim)
|
568 |
+
)
|
569 |
+
|
570 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
571 |
+
x = x.flatten(1) # Flatten spatial dimensions
|
572 |
+
return self.layers(x)
|
573 |
+
|
574 |
+
class SiameseSignatureNetwork(nn.Module):
|
575 |
+
"""Advanced Siamese network with precision-focused loss."""
|
576 |
+
|
577 |
+
def __init__(self, config: TrainingConfig = CONFIG):
|
578 |
+
super().__init__()
|
579 |
+
self.config = config
|
580 |
+
|
581 |
+
# Initialize backbone
|
582 |
+
if config.model_name.startswith("resnet"):
|
583 |
+
self.backbone = ResNetBackbone(
|
584 |
+
model_name=config.model_name,
|
585 |
+
pretrained_path=config.pretrained_path if config.last_epoch_weights is None else None
|
586 |
+
)
|
587 |
+
backbone_dim = self.backbone.output_dim
|
588 |
+
else:
|
589 |
+
raise ValueError(f"Unsupported model: {config.model_name}")
|
590 |
+
|
591 |
+
# Initialize embedding head
|
592 |
+
self.embedding_head = AdvancedEmbeddingHead(
|
593 |
+
input_dim=backbone_dim,
|
594 |
+
embedding_dim=config.embedding_dim,
|
595 |
+
dropout=0.5
|
596 |
+
)
|
597 |
+
|
598 |
+
self.normalize_embeddings = config.normalize_embeddings
|
599 |
+
self.distance_threshold = 0.5 # Will be updated during validation
|
600 |
+
|
601 |
+
# Loss components
|
602 |
+
self.criterion = MultiSimilarityLoss(
|
603 |
+
alpha=config.multisim_alpha,
|
604 |
+
beta=config.multisim_beta,
|
605 |
+
base=config.multisim_base
|
606 |
+
)
|
607 |
+
|
608 |
+
# Add triplet margin loss for better separation
|
609 |
+
self.triplet_loss = nn.TripletMarginLoss(
|
610 |
+
margin=config.triplet_margin,
|
611 |
+
p=2,
|
612 |
+
reduction='none' # We'll apply weights manually
|
613 |
+
)
|
614 |
+
|
615 |
+
# Loss weights
|
616 |
+
self.triplet_weight = config.triplet_weight
|
617 |
+
self.fp_penalty_weight = config.false_positive_penalty_weight
|
618 |
+
|
619 |
+
# Mining
|
620 |
+
if config.use_hard_mining:
|
621 |
+
self.miner = BatchHardMiner()
|
622 |
+
else:
|
623 |
+
self.miner = None
|
624 |
+
|
625 |
+
def get_parameter_groups(self) -> List[Dict[str, Any]]:
|
626 |
+
"""Get parameter groups for differential learning rates."""
|
627 |
+
backbone_params = list(self.backbone.parameters())
|
628 |
+
head_params = list(self.embedding_head.parameters())
|
629 |
+
|
630 |
+
return [
|
631 |
+
{
|
632 |
+
'params': backbone_params,
|
633 |
+
'lr': self.config.backbone_lr,
|
634 |
+
'name': 'backbone',
|
635 |
+
'weight_decay': self.config.weight_decay
|
636 |
+
},
|
637 |
+
{
|
638 |
+
'params': head_params,
|
639 |
+
'lr': self.config.head_lr,
|
640 |
+
'name': 'embedding_head',
|
641 |
+
'weight_decay': self.config.weight_decay
|
642 |
+
}
|
643 |
+
]
|
644 |
+
|
645 |
+
def forward(self, anchor: torch.Tensor, positive: torch.Tensor,
|
646 |
+
negative: Optional[torch.Tensor] = None) -> Union[Tuple[torch.Tensor, torch.Tensor],
|
647 |
+
Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
|
648 |
+
"""Forward pass for training or inference."""
|
649 |
+
a_features = self.backbone(anchor)
|
650 |
+
a_emb = self.embedding_head(a_features)
|
651 |
+
|
652 |
+
p_features = self.backbone(positive)
|
653 |
+
p_emb = self.embedding_head(p_features)
|
654 |
+
|
655 |
+
if self.normalize_embeddings:
|
656 |
+
a_emb = F.normalize(a_emb, p=2, dim=1)
|
657 |
+
p_emb = F.normalize(p_emb, p=2, dim=1)
|
658 |
+
|
659 |
+
if negative is not None:
|
660 |
+
n_features = self.backbone(negative)
|
661 |
+
n_emb = self.embedding_head(n_features)
|
662 |
+
|
663 |
+
if self.normalize_embeddings:
|
664 |
+
n_emb = F.normalize(n_emb, p=2, dim=1)
|
665 |
+
|
666 |
+
return a_emb, p_emb, n_emb
|
667 |
+
|
668 |
+
return a_emb, p_emb
|
669 |
+
|
670 |
+
def compute_loss(self, embeddings: torch.Tensor, labels: torch.Tensor,
|
671 |
+
anchors: Optional[torch.Tensor] = None,
|
672 |
+
positives: Optional[torch.Tensor] = None,
|
673 |
+
negatives: Optional[torch.Tensor] = None,
|
674 |
+
distance_weights: Optional[Dict[str, torch.Tensor]] = None) -> torch.Tensor:
|
675 |
+
"""Enhanced loss computation with precision focus and distance weighting."""
|
676 |
+
|
677 |
+
# MultiSimilarity loss
|
678 |
+
if self.miner is not None:
|
679 |
+
hard_pairs = self.miner(embeddings, labels)
|
680 |
+
ms_loss = self.criterion(embeddings, labels, hard_pairs)
|
681 |
+
else:
|
682 |
+
ms_loss = self.criterion(embeddings, labels)
|
683 |
+
|
684 |
+
total_loss = ms_loss
|
685 |
+
|
686 |
+
# Add triplet loss if embeddings provided
|
687 |
+
if anchors is not None and positives is not None and negatives is not None:
|
688 |
+
# Compute triplet losses for each sample
|
689 |
+
triplet_losses = self.triplet_loss(anchors, positives, negatives)
|
690 |
+
|
691 |
+
# Apply distance-based weights if provided
|
692 |
+
if distance_weights is not None:
|
693 |
+
neg_weights = distance_weights.get('negative_weights', torch.ones_like(triplet_losses))
|
694 |
+
weighted_triplet_loss = (triplet_losses * neg_weights).mean()
|
695 |
+
else:
|
696 |
+
weighted_triplet_loss = triplet_losses.mean()
|
697 |
+
|
698 |
+
total_loss += self.triplet_weight * weighted_triplet_loss
|
699 |
+
|
700 |
+
# Additional penalty for hard negatives (false positives)
|
701 |
+
with torch.no_grad():
|
702 |
+
d_an = F.pairwise_distance(anchors, negatives)
|
703 |
+
# Find negatives that are too close (potential false positives)
|
704 |
+
hard_negative_mask = d_an < self.distance_threshold
|
705 |
+
|
706 |
+
if hard_negative_mask.any():
|
707 |
+
# Apply distance-based weights for false positive penalty
|
708 |
+
if distance_weights is not None:
|
709 |
+
neg_weights = distance_weights.get('negative_weights', torch.ones_like(d_an))
|
710 |
+
# Extra penalty weighted by how bad the false positive is
|
711 |
+
false_positive_distances = self.distance_threshold - d_an[hard_negative_mask]
|
712 |
+
false_positive_weights = neg_weights[hard_negative_mask]
|
713 |
+
fp_loss = (false_positive_distances * false_positive_weights).mean()
|
714 |
+
else:
|
715 |
+
fp_loss = (self.distance_threshold - d_an[hard_negative_mask]).mean()
|
716 |
+
|
717 |
+
total_loss += self.fp_penalty_weight * fp_loss
|
718 |
+
|
719 |
+
return total_loss
|
720 |
+
|
721 |
+
def predict_pair(self, img1: torch.Tensor, img2: torch.Tensor,
|
722 |
+
threshold: Optional[float] = None, return_dist: bool = False) -> torch.Tensor:
|
723 |
+
"""Predict similarity between image pairs."""
|
724 |
+
self.eval()
|
725 |
+
with torch.no_grad():
|
726 |
+
emb1, emb2 = self(img1, img2)
|
727 |
+
distances = F.pairwise_distance(emb1, emb2)
|
728 |
+
|
729 |
+
if return_dist:
|
730 |
+
return distances
|
731 |
+
|
732 |
+
thresh = threshold if threshold is not None else self.distance_threshold
|
733 |
+
return (distances < thresh).long()
|
734 |
+
|
735 |
+
# ----------------------------
|
736 |
+
# Advanced Training Metrics and Statistics
|
737 |
+
# ----------------------------
|
738 |
+
class TrainingMetrics:
|
739 |
+
"""Enhanced training metrics with adaptive mining for both hard positives and negatives."""
|
740 |
+
|
741 |
+
def __init__(self):
|
742 |
+
self.reset()
|
743 |
+
# Track distance statistics for adaptive thresholds
|
744 |
+
self.distance_history = {"positive": [], "negative": []}
|
745 |
+
self.adaptive_stats = {}
|
746 |
+
|
747 |
+
def reset(self):
|
748 |
+
"""Reset all metrics."""
|
749 |
+
self.losses = []
|
750 |
+
self.genuine_distances = []
|
751 |
+
self.forged_distances = []
|
752 |
+
|
753 |
+
# Separate mining stats for positives and negatives
|
754 |
+
self.positive_mining_stats = {"easy": 0, "semi": 0, "hard": 0}
|
755 |
+
self.negative_mining_stats = {"easy": 0, "semi": 0, "hard": 0}
|
756 |
+
|
757 |
+
# Hard sample counts
|
758 |
+
self.hard_positive_count = 0
|
759 |
+
self.hard_negative_count = 0
|
760 |
+
self.total_positive_pairs = 0
|
761 |
+
self.total_negative_pairs = 0
|
762 |
+
|
763 |
+
# False positive/negative tracking
|
764 |
+
self.false_positive_count = 0
|
765 |
+
self.false_negative_count = 0
|
766 |
+
|
767 |
+
self.learning_rates = {}
|
768 |
+
|
769 |
+
def update_distance_statistics(self, d_positive: np.ndarray, d_negative: np.ndarray):
|
770 |
+
"""Update running statistics for adaptive thresholds."""
|
771 |
+
# Keep rolling window of recent distances
|
772 |
+
self.distance_history["positive"].extend(d_positive.tolist())
|
773 |
+
self.distance_history["negative"].extend(d_negative.tolist())
|
774 |
+
|
775 |
+
# Keep only recent history (last 5000 samples)
|
776 |
+
for key in self.distance_history:
|
777 |
+
if len(self.distance_history[key]) > 5000:
|
778 |
+
self.distance_history[key] = self.distance_history[key][-5000:]
|
779 |
+
|
780 |
+
# Compute adaptive statistics
|
781 |
+
if len(self.distance_history["positive"]) > 100 and len(self.distance_history["negative"]) > 100:
|
782 |
+
pos_distances = np.array(self.distance_history["positive"])
|
783 |
+
neg_distances = np.array(self.distance_history["negative"])
|
784 |
+
|
785 |
+
self.adaptive_stats = {
|
786 |
+
"pos_mean": np.mean(pos_distances),
|
787 |
+
"pos_std": np.std(pos_distances),
|
788 |
+
"pos_q25": np.percentile(pos_distances, 25),
|
789 |
+
"pos_q50": np.percentile(pos_distances, 50),
|
790 |
+
"pos_q75": np.percentile(pos_distances, 75),
|
791 |
+
"pos_q90": np.percentile(pos_distances, 90),
|
792 |
+
|
793 |
+
"neg_mean": np.mean(neg_distances),
|
794 |
+
"neg_std": np.std(neg_distances),
|
795 |
+
"neg_q10": np.percentile(neg_distances, 10),
|
796 |
+
"neg_q25": np.percentile(neg_distances, 25),
|
797 |
+
"neg_q50": np.percentile(neg_distances, 50),
|
798 |
+
"neg_q75": np.percentile(neg_distances, 75),
|
799 |
+
|
800 |
+
"separation": np.mean(neg_distances) - np.mean(pos_distances),
|
801 |
+
"overlap_region": max(0, np.percentile(pos_distances, 95) - np.percentile(neg_distances, 5))
|
802 |
+
}
|
803 |
+
|
804 |
+
def compute_precision_focused_weights(self, d_positive: np.ndarray,
|
805 |
+
d_negative: np.ndarray,
|
806 |
+
negative_weight_multiplier: float = 2.5) -> Tuple[torch.Tensor, torch.Tensor]:
|
807 |
+
"""Compute sample weights with focus on improving precision."""
|
808 |
+
pos_weights = np.ones_like(d_positive)
|
809 |
+
neg_weights = np.ones_like(d_negative)
|
810 |
+
|
811 |
+
if self.adaptive_stats:
|
812 |
+
# Hard negatives (forged that look genuine) get MUCH higher weight
|
813 |
+
neg_q10 = self.adaptive_stats["neg_q10"]
|
814 |
+
neg_q25 = self.adaptive_stats["neg_q25"]
|
815 |
+
|
816 |
+
# Very hard negatives (bottom 10%) - highest weight
|
817 |
+
very_hard_neg_mask = d_negative < neg_q10
|
818 |
+
neg_weights[very_hard_neg_mask] = negative_weight_multiplier * 1.5
|
819 |
+
|
820 |
+
# Hard negatives (10-25%) - high weight
|
821 |
+
hard_neg_mask = (d_negative >= neg_q10) & (d_negative < neg_q25)
|
822 |
+
neg_weights[hard_neg_mask] = negative_weight_multiplier
|
823 |
+
|
824 |
+
# Semi-hard negatives (25-50%) - moderate weight
|
825 |
+
semi_neg_mask = (d_negative >= neg_q25) & (d_negative < self.adaptive_stats["neg_q50"])
|
826 |
+
neg_weights[semi_neg_mask] = negative_weight_multiplier * 0.6
|
827 |
+
|
828 |
+
# Hard positives get moderate weight (but less than hard negatives)
|
829 |
+
pos_q75 = self.adaptive_stats["pos_q75"]
|
830 |
+
pos_q90 = self.adaptive_stats["pos_q90"]
|
831 |
+
|
832 |
+
# Very hard positives (top 10%)
|
833 |
+
very_hard_pos_mask = d_positive > pos_q90
|
834 |
+
pos_weights[very_hard_pos_mask] = 1.8
|
835 |
+
|
836 |
+
# Hard positives (75-90%)
|
837 |
+
hard_pos_mask = (d_positive > pos_q75) & (d_positive <= pos_q90)
|
838 |
+
pos_weights[hard_pos_mask] = 1.5
|
839 |
+
|
840 |
+
return torch.tensor(pos_weights, dtype=torch.float32), torch.tensor(neg_weights, dtype=torch.float32)
|
841 |
+
|
842 |
+
def update_mining_stats(self, d_positive: np.ndarray, d_negative: np.ndarray,
|
843 |
+
margin: float, difficulty_params: Dict[str, float]):
|
844 |
+
"""Intelligent adaptive mining for both hard positives and hard negatives."""
|
845 |
+
|
846 |
+
# Update distance statistics first
|
847 |
+
self.update_distance_statistics(d_positive, d_negative)
|
848 |
+
|
849 |
+
# Update totals
|
850 |
+
self.total_positive_pairs += len(d_positive)
|
851 |
+
self.total_negative_pairs += len(d_negative)
|
852 |
+
|
853 |
+
# Use adaptive thresholds if available, otherwise fallback to fixed
|
854 |
+
if self.adaptive_stats:
|
855 |
+
self._adaptive_mining(d_positive, d_negative, difficulty_params)
|
856 |
+
else:
|
857 |
+
self._fixed_mining(d_positive, d_negative, margin)
|
858 |
+
|
859 |
+
def _adaptive_mining(self, d_positive: np.ndarray, d_negative: np.ndarray,
|
860 |
+
difficulty_params: Dict[str, float]):
|
861 |
+
"""Adaptive mining based on current distance distributions."""
|
862 |
+
stats = self.adaptive_stats
|
863 |
+
|
864 |
+
# Get difficulty parameters
|
865 |
+
hard_positive_ratio = difficulty_params.get("hard_positive_ratio", 0.3)
|
866 |
+
hard_negative_ratio = difficulty_params.get("hard_negative_ratio", 0.3)
|
867 |
+
|
868 |
+
# Dynamic thresholds for hard positives (far apart genuine pairs)
|
869 |
+
# Use percentile based on desired hard positive ratio
|
870 |
+
hard_pos_percentile = 100 - (hard_positive_ratio * 100)
|
871 |
+
hard_pos_threshold = np.percentile(self.distance_history["positive"][-1000:], hard_pos_percentile)
|
872 |
+
semi_pos_threshold = stats["pos_q50"]
|
873 |
+
|
874 |
+
# Dynamic thresholds for hard negatives (close together impostor pairs)
|
875 |
+
# Use percentile based on desired hard negative ratio
|
876 |
+
hard_neg_percentile = hard_negative_ratio * 100
|
877 |
+
hard_neg_threshold = np.percentile(self.distance_history["negative"][-1000:], hard_neg_percentile)
|
878 |
+
semi_neg_threshold = stats["neg_q50"]
|
879 |
+
|
880 |
+
# Mine hard positives
|
881 |
+
for dp in d_positive:
|
882 |
+
if dp >= hard_pos_threshold:
|
883 |
+
self.positive_mining_stats["hard"] += 1
|
884 |
+
self.hard_positive_count += 1
|
885 |
+
elif dp >= semi_pos_threshold:
|
886 |
+
self.positive_mining_stats["semi"] += 1
|
887 |
+
else:
|
888 |
+
self.positive_mining_stats["easy"] += 1
|
889 |
+
|
890 |
+
# Mine hard negatives
|
891 |
+
for dn in d_negative:
|
892 |
+
if dn <= hard_neg_threshold:
|
893 |
+
self.negative_mining_stats["hard"] += 1
|
894 |
+
self.hard_negative_count += 1
|
895 |
+
elif dn <= semi_neg_threshold:
|
896 |
+
self.negative_mining_stats["semi"] += 1
|
897 |
+
else:
|
898 |
+
self.negative_mining_stats["easy"] += 1
|
899 |
+
|
900 |
+
def _fixed_mining(self, d_positive: np.ndarray, d_negative: np.ndarray, margin: float):
|
901 |
+
"""Fallback fixed mining for early epochs."""
|
902 |
+
# Fixed thresholds
|
903 |
+
hard_pos_threshold = 0.5 # Far genuine pairs
|
904 |
+
hard_neg_threshold = 0.3 # Close impostor pairs
|
905 |
+
|
906 |
+
for dp in d_positive:
|
907 |
+
if dp >= hard_pos_threshold:
|
908 |
+
self.positive_mining_stats["hard"] += 1
|
909 |
+
self.hard_positive_count += 1
|
910 |
+
elif dp >= hard_pos_threshold * 0.7:
|
911 |
+
self.positive_mining_stats["semi"] += 1
|
912 |
+
else:
|
913 |
+
self.positive_mining_stats["easy"] += 1
|
914 |
+
|
915 |
+
for dn in d_negative:
|
916 |
+
if dn <= hard_neg_threshold:
|
917 |
+
self.negative_mining_stats["hard"] += 1
|
918 |
+
self.hard_negative_count += 1
|
919 |
+
elif dn <= hard_neg_threshold * 1.5:
|
920 |
+
self.negative_mining_stats["semi"] += 1
|
921 |
+
else:
|
922 |
+
self.negative_mining_stats["easy"] += 1
|
923 |
+
|
924 |
+
def get_mining_percentages(self) -> Dict[str, float]:
|
925 |
+
"""Get mining statistics as percentages with debugging info."""
|
926 |
+
total_pos = sum(self.positive_mining_stats.values())
|
927 |
+
total_neg = sum(self.negative_mining_stats.values())
|
928 |
+
|
929 |
+
percentages = {}
|
930 |
+
|
931 |
+
# Positive pair mining stats
|
932 |
+
if total_pos > 0:
|
933 |
+
percentages.update({
|
934 |
+
"pos_mining_easy_pct": 100.0 * self.positive_mining_stats["easy"] / total_pos,
|
935 |
+
"pos_mining_semi_pct": 100.0 * self.positive_mining_stats["semi"] / total_pos,
|
936 |
+
"pos_mining_hard_pct": 100.0 * self.positive_mining_stats["hard"] / total_pos,
|
937 |
+
})
|
938 |
+
else:
|
939 |
+
percentages.update({
|
940 |
+
"pos_mining_easy_pct": 0.0,
|
941 |
+
"pos_mining_semi_pct": 0.0,
|
942 |
+
"pos_mining_hard_pct": 0.0,
|
943 |
+
})
|
944 |
+
|
945 |
+
# Negative pair mining stats
|
946 |
+
if total_neg > 0:
|
947 |
+
percentages.update({
|
948 |
+
"neg_mining_easy_pct": 100.0 * self.negative_mining_stats["easy"] / total_neg,
|
949 |
+
"neg_mining_semi_pct": 100.0 * self.negative_mining_stats["semi"] / total_neg,
|
950 |
+
"neg_mining_hard_pct": 100.0 * self.negative_mining_stats["hard"] / total_neg,
|
951 |
+
})
|
952 |
+
else:
|
953 |
+
percentages.update({
|
954 |
+
"neg_mining_easy_pct": 0.0,
|
955 |
+
"neg_mining_semi_pct": 0.0,
|
956 |
+
"neg_mining_hard_pct": 0.0,
|
957 |
+
})
|
958 |
+
|
959 |
+
# Overall hard sample ratios
|
960 |
+
if self.total_positive_pairs > 0:
|
961 |
+
percentages["hard_positive_ratio"] = 100.0 * self.hard_positive_count / self.total_positive_pairs
|
962 |
+
else:
|
963 |
+
percentages["hard_positive_ratio"] = 0.0
|
964 |
+
|
965 |
+
if self.total_negative_pairs > 0:
|
966 |
+
percentages["hard_negative_ratio"] = 100.0 * self.hard_negative_count / self.total_negative_pairs
|
967 |
+
else:
|
968 |
+
percentages["hard_negative_ratio"] = 0.0
|
969 |
+
|
970 |
+
# False positive/negative rates
|
971 |
+
total_samples = self.total_positive_pairs + self.total_negative_pairs
|
972 |
+
if total_samples > 0:
|
973 |
+
percentages["false_positive_rate"] = 100.0 * self.false_positive_count / self.total_negative_pairs if self.total_negative_pairs > 0 else 0.0
|
974 |
+
percentages["false_negative_rate"] = 100.0 * self.false_negative_count / self.total_positive_pairs if self.total_positive_pairs > 0 else 0.0
|
975 |
+
|
976 |
+
# Add adaptive stats if available
|
977 |
+
if self.adaptive_stats:
|
978 |
+
percentages.update({
|
979 |
+
"adaptive_separation": self.adaptive_stats["separation"],
|
980 |
+
"adaptive_overlap": self.adaptive_stats["overlap_region"],
|
981 |
+
"adaptive_pos_spread": self.adaptive_stats["pos_std"],
|
982 |
+
"adaptive_neg_spread": self.adaptive_stats["neg_std"],
|
983 |
+
})
|
984 |
+
|
985 |
+
return percentages
|
986 |
+
|
987 |
+
def compute_separation_metrics(self) -> Dict[str, float]:
|
988 |
+
"""Compute distance separation metrics."""
|
989 |
+
if not self.genuine_distances or not self.forged_distances:
|
990 |
+
return {
|
991 |
+
"genuine_dist_mean": 0.0,
|
992 |
+
"forged_dist_mean": 0.0,
|
993 |
+
"genuine_dist_std": 0.0,
|
994 |
+
"forged_dist_std": 0.0,
|
995 |
+
"separation": 0.0,
|
996 |
+
"overlap": 0.0,
|
997 |
+
"separation_ratio": 0.0,
|
998 |
+
"cohesion_ratio": 0.0
|
999 |
+
}
|
1000 |
+
|
1001 |
+
gen_mean = np.mean(self.genuine_distances)
|
1002 |
+
forg_mean = np.mean(self.forged_distances)
|
1003 |
+
gen_std = np.std(self.genuine_distances)
|
1004 |
+
forg_std = np.std(self.forged_distances)
|
1005 |
+
|
1006 |
+
separation = forg_mean - gen_mean
|
1007 |
+
overlap = max(0, gen_mean + 2*gen_std - (forg_mean - 2*forg_std))
|
1008 |
+
|
1009 |
+
# Cohesion ratio: how tight are genuine pairs relative to separation
|
1010 |
+
cohesion_ratio = gen_std / (separation + 1e-8)
|
1011 |
+
|
1012 |
+
return {
|
1013 |
+
"genuine_dist_mean": gen_mean,
|
1014 |
+
"forged_dist_mean": forg_mean,
|
1015 |
+
"genuine_dist_std": gen_std,
|
1016 |
+
"forged_dist_std": forg_std,
|
1017 |
+
"separation": separation,
|
1018 |
+
"overlap": overlap,
|
1019 |
+
"separation_ratio": separation / (gen_std + forg_std + 1e-8),
|
1020 |
+
"cohesion_ratio": cohesion_ratio
|
1021 |
+
}
|
1022 |
+
|
1023 |
+
# ----------------------------
|
1024 |
+
# Enhanced Training Loop
|
1025 |
+
# ----------------------------
|
1026 |
+
class SignatureTrainer:
|
1027 |
+
"""Research-grade signature verification trainer."""
|
1028 |
+
|
1029 |
+
def __init__(self, config: TrainingConfig = CONFIG):
|
1030 |
+
self.config = config
|
1031 |
+
self.device = torch.device(config.device)
|
1032 |
+
|
1033 |
+
# Initialize managers
|
1034 |
+
self.mlflow_manager = MLFlowManager(tracking_uri=self.config.tracking_uri)
|
1035 |
+
self.curriculum_manager = CurriculumLearningManager(config)
|
1036 |
+
|
1037 |
+
# Training state
|
1038 |
+
self.current_epoch = 0
|
1039 |
+
self.best_eer = float('inf')
|
1040 |
+
self.patience_counter = 0
|
1041 |
+
self.global_step = 0
|
1042 |
+
|
1043 |
+
# Setup logging
|
1044 |
+
self._setup_logging()
|
1045 |
+
|
1046 |
+
def _setup_logging(self):
|
1047 |
+
"""Setup comprehensive logging."""
|
1048 |
+
logging.basicConfig(
|
1049 |
+
level=logging.INFO,
|
1050 |
+
format='%(asctime)s - %(levelname)s - %(message)s',
|
1051 |
+
handlers=[
|
1052 |
+
logging.FileHandler('training.log'),
|
1053 |
+
logging.StreamHandler()
|
1054 |
+
]
|
1055 |
+
)
|
1056 |
+
self.logger = logging.getLogger(__name__)
|
1057 |
+
|
1058 |
+
def _prepare_datasets(self) -> Tuple[SignatureDataset, SignatureDataset]:
|
1059 |
+
"""Prepare training and validation datasets."""
|
1060 |
+
# Load datasets
|
1061 |
+
train_data = pd.read_excel("../../data/classify/preprared_data/labels/train_triplets_balanced_v12.xlsx")
|
1062 |
+
val_data = pd.read_excel("../../data/classify/preprared_data/labels/valid_pairs_balanced_v12.xlsx")
|
1063 |
+
|
1064 |
+
train_dataset = SignatureDataset(
|
1065 |
+
folder_img="../../data/classify/preprared_data/images/",
|
1066 |
+
excel_data=train_data,
|
1067 |
+
curriculum_manager=self.curriculum_manager,
|
1068 |
+
is_train=True,
|
1069 |
+
config=self.config
|
1070 |
+
)
|
1071 |
+
|
1072 |
+
val_dataset = SignatureDataset(
|
1073 |
+
folder_img="../../data/classify/preprared_data/images/",
|
1074 |
+
excel_data=val_data,
|
1075 |
+
curriculum_manager=self.curriculum_manager,
|
1076 |
+
is_train=False,
|
1077 |
+
config=self.config
|
1078 |
+
)
|
1079 |
+
|
1080 |
+
self.logger.info(f"Training samples: {len(train_dataset)}")
|
1081 |
+
self.logger.info(f"Validation samples: {len(val_dataset)}")
|
1082 |
+
|
1083 |
+
return train_dataset, val_dataset
|
1084 |
+
|
1085 |
+
def _compute_precision_optimized_threshold(self, distances: np.ndarray,
|
1086 |
+
labels: np.ndarray,
|
1087 |
+
target_precision: float = None) -> float:
|
1088 |
+
"""Find threshold that achieves target precision while maximizing F1."""
|
1089 |
+
if target_precision is None:
|
1090 |
+
target_precision = self.config.target_precision
|
1091 |
+
|
1092 |
+
thresholds = np.linspace(distances.min(), distances.max(), 1000)
|
1093 |
+
best_threshold = thresholds[0]
|
1094 |
+
best_f1 = 0
|
1095 |
+
best_precision = 0
|
1096 |
+
best_recall = 0
|
1097 |
+
|
1098 |
+
for thresh in thresholds:
|
1099 |
+
predictions = (distances < thresh).astype(int)
|
1100 |
+
|
1101 |
+
# Calculate metrics
|
1102 |
+
tp = np.sum((predictions == 1) & (labels == 1))
|
1103 |
+
fp = np.sum((predictions == 1) & (labels == 0))
|
1104 |
+
fn = np.sum((predictions == 0) & (labels == 1))
|
1105 |
+
|
1106 |
+
precision = tp / (tp + fp + 1e-8)
|
1107 |
+
recall = tp / (tp + fn + 1e-8)
|
1108 |
+
f1 = 2 * precision * recall / (precision + recall + 1e-8)
|
1109 |
+
|
1110 |
+
# Prioritize precision while maintaining reasonable recall
|
1111 |
+
if precision >= target_precision and f1 > best_f1:
|
1112 |
+
best_f1 = f1
|
1113 |
+
best_threshold = thresh
|
1114 |
+
best_precision = precision
|
1115 |
+
best_recall = recall
|
1116 |
+
# If we can't achieve target precision, get best precision with recall > 0.5
|
1117 |
+
elif precision > best_precision and recall > 0.5:
|
1118 |
+
best_f1 = f1
|
1119 |
+
best_threshold = thresh
|
1120 |
+
best_precision = precision
|
1121 |
+
best_recall = recall
|
1122 |
+
|
1123 |
+
print(f" Precision-optimized threshold: {best_threshold:.4f} "
|
1124 |
+
f"(P: {best_precision:.3f}, R: {best_recall:.3f}, F1: {best_f1:.3f})")
|
1125 |
+
|
1126 |
+
return best_threshold
|
1127 |
+
|
1128 |
+
def _setup_model_and_optimizer(self) -> Tuple[SiameseSignatureNetwork, torch.optim.Optimizer, Any]:
|
1129 |
+
"""Setup model, optimizer, and scheduler."""
|
1130 |
+
# Initialize model
|
1131 |
+
model = SiameseSignatureNetwork(self.config)
|
1132 |
+
|
1133 |
+
# Compile model if available
|
1134 |
+
if hasattr(torch, "compile") and platform.system() != "Windows":
|
1135 |
+
self.logger.info("Compiling model with torch.compile")
|
1136 |
+
model = torch.compile(model)
|
1137 |
+
|
1138 |
+
model = model.to(self.device)
|
1139 |
+
|
1140 |
+
# Count parameters
|
1141 |
+
total_params = sum(p.numel() for p in model.parameters())
|
1142 |
+
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
1143 |
+
self.logger.info(f"Total parameters: {total_params:,}")
|
1144 |
+
self.logger.info(f"Trainable parameters: {trainable_params:,}")
|
1145 |
+
|
1146 |
+
# Setup optimizer with parameter groups
|
1147 |
+
param_groups = model.get_parameter_groups()
|
1148 |
+
optimizer = torch.optim.AdamW(param_groups)
|
1149 |
+
|
1150 |
+
# Log learning rates
|
1151 |
+
for group in param_groups:
|
1152 |
+
self.logger.info(f"Parameter group '{group['name']}': LR = {group['lr']:.2e}")
|
1153 |
+
|
1154 |
+
# Setup scheduler
|
1155 |
+
if self.config.lr_scheduler == "cosine":
|
1156 |
+
scheduler = CosineAnnealingLR(optimizer, T_max=self.config.max_epochs)
|
1157 |
+
else:
|
1158 |
+
scheduler = ReduceLROnPlateau(optimizer, mode='min', patience=5, factor=0.5)
|
1159 |
+
|
1160 |
+
return model, optimizer, scheduler
|
1161 |
+
|
1162 |
+
def _setup_checkpoint_management(self, run_id: str) -> Tuple[str, str]:
|
1163 |
+
"""Setup checkpoint directories."""
|
1164 |
+
checkpoint_dir = os.path.join("../../model/models_checkpoints/", run_id)
|
1165 |
+
figures_dir = os.path.join(checkpoint_dir, "figures")
|
1166 |
+
os.makedirs(checkpoint_dir, exist_ok=True)
|
1167 |
+
os.makedirs(figures_dir, exist_ok=True)
|
1168 |
+
return checkpoint_dir, figures_dir
|
1169 |
+
|
1170 |
+
def _load_checkpoint(self, model: nn.Module, optimizer: torch.optim.Optimizer,
|
1171 |
+
scheduler: Any, scaler: torch.amp.GradScaler) -> int:
|
1172 |
+
"""Load checkpoint if specified."""
|
1173 |
+
if not self.config.last_epoch_weights:
|
1174 |
+
return 1
|
1175 |
+
|
1176 |
+
checkpoint_path = self.config.last_epoch_weights
|
1177 |
+
self.logger.info(f"Loading checkpoint from {checkpoint_path}")
|
1178 |
+
|
1179 |
+
try:
|
1180 |
+
checkpoint = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
|
1181 |
+
|
1182 |
+
model.load_state_dict(checkpoint["model_state_dict"])
|
1183 |
+
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
|
1184 |
+
scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
|
1185 |
+
scaler.load_state_dict(checkpoint.get("scaler_state_dict", scaler.state_dict()))
|
1186 |
+
|
1187 |
+
start_epoch = checkpoint["epoch"] + 1
|
1188 |
+
self.best_eer = checkpoint.get("best_eer", self.best_eer)
|
1189 |
+
model.distance_threshold = checkpoint.get("prediction_threshold", 0.5)
|
1190 |
+
|
1191 |
+
self.logger.info(f"Resumed from epoch {start_epoch}, best EER: {self.best_eer:.4f}")
|
1192 |
+
return start_epoch
|
1193 |
+
|
1194 |
+
except Exception as e:
|
1195 |
+
self.logger.error(f"Failed to load checkpoint: {e}")
|
1196 |
+
return 1
|
1197 |
+
|
1198 |
+
def train_epoch(self, model: nn.Module, train_loader: DataLoader,
|
1199 |
+
optimizer: torch.optim.Optimizer, scaler: torch.amp.GradScaler,
|
1200 |
+
epoch: int) -> TrainingMetrics:
|
1201 |
+
"""Enhanced training with intelligent adaptive mining for both hard positives and negatives."""
|
1202 |
+
model.train()
|
1203 |
+
metrics = TrainingMetrics()
|
1204 |
+
|
1205 |
+
curriculum_stats = train_loader.dataset.get_curriculum_stats()
|
1206 |
+
|
1207 |
+
# INTELLIGENT MARGIN CALCULATION
|
1208 |
+
base_margin = 0.5 # Base margin for normalized embeddings
|
1209 |
+
margin_multiplier = curriculum_stats["margin_multiplier"]
|
1210 |
+
adaptive_margin = base_margin * margin_multiplier
|
1211 |
+
|
1212 |
+
# Progressive margin adjustment based on epoch
|
1213 |
+
epoch_progress = epoch / self.config.max_epochs
|
1214 |
+
progressive_factor = 1.2 - 0.4 * epoch_progress # 1.2 → 0.8
|
1215 |
+
final_margin = adaptive_margin * progressive_factor
|
1216 |
+
|
1217 |
+
# Tracking counters
|
1218 |
+
forgery_batch_count = 0
|
1219 |
+
genuine_batch_count = 0
|
1220 |
+
batch_fp_count = 0
|
1221 |
+
batch_fn_count = 0
|
1222 |
+
|
1223 |
+
# Debug info
|
1224 |
+
debug_printed = False
|
1225 |
+
|
1226 |
+
pbar = tqdm(train_loader, desc=f"[Train] Epoch {epoch}")
|
1227 |
+
|
1228 |
+
for step, (anchors, positives, negatives, anchor_ids, negative_ids) in enumerate(pbar):
|
1229 |
+
|
1230 |
+
# Move to device
|
1231 |
+
anchors = anchors.to(self.device, non_blocking=True)
|
1232 |
+
positives = positives.to(self.device, non_blocking=True)
|
1233 |
+
negatives = negatives.to(self.device, non_blocking=True)
|
1234 |
+
anchor_ids = anchor_ids.to(self.device, non_blocking=True)
|
1235 |
+
negative_ids = negative_ids.to(self.device, non_blocking=True)
|
1236 |
+
|
1237 |
+
# Count forgery vs genuine negatives
|
1238 |
+
max_genuine_id = len(train_loader.dataset.genuine_id_mapping)
|
1239 |
+
forgery_mask = negative_ids >= max_genuine_id + 100
|
1240 |
+
forgery_batch_count += forgery_mask.sum().item()
|
1241 |
+
genuine_batch_count += (~forgery_mask).sum().item()
|
1242 |
+
|
1243 |
+
if not debug_printed and step == 0:
|
1244 |
+
print(f"\n[DEBUG Epoch {epoch}]")
|
1245 |
+
print(f" Final margin: {final_margin:.3f}")
|
1246 |
+
print(f" Hard negative ratio target: {curriculum_stats['hard_negative_ratio']:.3f}")
|
1247 |
+
print(f" Hard positive ratio target: {curriculum_stats['hard_positive_ratio']:.3f}")
|
1248 |
+
print(f" Negative weight multiplier: {self.config.negative_weight_multiplier:.2f}")
|
1249 |
+
print(f" Triplet weight: {self.config.triplet_weight:.2f}")
|
1250 |
+
print(f" FP penalty weight: {self.config.false_positive_penalty_weight:.2f}")
|
1251 |
+
debug_printed = True
|
1252 |
+
|
1253 |
+
# Forward pass to get embeddings first
|
1254 |
+
with torch.amp.autocast(device_type=self.device.type):
|
1255 |
+
a_emb, p_emb, n_emb = model(anchors, positives, negatives)
|
1256 |
+
|
1257 |
+
# Compute distances and weights BEFORE loss computation
|
1258 |
+
with torch.no_grad():
|
1259 |
+
d_ap = F.pairwise_distance(a_emb, p_emb).cpu().numpy()
|
1260 |
+
d_an = F.pairwise_distance(a_emb, n_emb).cpu().numpy()
|
1261 |
+
|
1262 |
+
# Get precision-focused weights
|
1263 |
+
pos_weights, neg_weights = metrics.compute_precision_focused_weights(
|
1264 |
+
d_ap, d_an,
|
1265 |
+
negative_weight_multiplier=self.config.negative_weight_multiplier
|
1266 |
+
)
|
1267 |
+
pos_weights = pos_weights.to(self.device)
|
1268 |
+
neg_weights = neg_weights.to(self.device)
|
1269 |
+
|
1270 |
+
# Track false positives/negatives
|
1271 |
+
fp_mask = d_an < model.distance_threshold
|
1272 |
+
fn_mask = d_ap > model.distance_threshold
|
1273 |
+
batch_fp_count = fp_mask.sum()
|
1274 |
+
batch_fn_count = fn_mask.sum()
|
1275 |
+
metrics.false_positive_count += batch_fp_count
|
1276 |
+
metrics.false_negative_count += batch_fn_count
|
1277 |
+
|
1278 |
+
# Prepare distance weights for loss
|
1279 |
+
distance_weights = {
|
1280 |
+
'positive_weights': pos_weights,
|
1281 |
+
'negative_weights': neg_weights
|
1282 |
+
}
|
1283 |
+
|
1284 |
+
# Now compute loss with weights
|
1285 |
+
with torch.amp.autocast(device_type=self.device.type):
|
1286 |
+
all_embeddings = torch.cat([a_emb, p_emb, n_emb], dim=0)
|
1287 |
+
all_labels = torch.cat([anchor_ids, anchor_ids, negative_ids], dim=0)
|
1288 |
+
|
1289 |
+
# Compute loss with triplet component and distance weights
|
1290 |
+
batch_loss = model.compute_loss(
|
1291 |
+
all_embeddings, all_labels,
|
1292 |
+
anchors=a_emb, positives=p_emb, negatives=n_emb,
|
1293 |
+
distance_weights=distance_weights
|
1294 |
+
)
|
1295 |
+
|
1296 |
+
# Gradient accumulation
|
1297 |
+
loss = batch_loss / self.config.grad_accum_steps
|
1298 |
+
scaler.scale(loss).backward()
|
1299 |
+
|
1300 |
+
if (step + 1) % self.config.grad_accum_steps == 0 or (step + 1) == len(train_loader):
|
1301 |
+
scaler.unscale_(optimizer)
|
1302 |
+
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
|
1303 |
+
scaler.step(optimizer)
|
1304 |
+
scaler.update()
|
1305 |
+
optimizer.zero_grad(set_to_none=True)
|
1306 |
+
self.global_step += 1
|
1307 |
+
|
1308 |
+
# Update metrics
|
1309 |
+
metrics.losses.append(batch_loss.item())
|
1310 |
+
metrics.genuine_distances.extend(d_ap.tolist())
|
1311 |
+
metrics.forged_distances.extend(d_an.tolist())
|
1312 |
+
|
1313 |
+
# Use enhanced mining with difficulty parameters
|
1314 |
+
metrics.update_mining_stats(d_ap, d_an, final_margin, curriculum_stats)
|
1315 |
+
|
1316 |
+
# Store learning rates
|
1317 |
+
for i, group in enumerate(optimizer.param_groups):
|
1318 |
+
metrics.learning_rates[f"lr_{group.get('name', i)}"] = group['lr']
|
1319 |
+
|
1320 |
+
# Enhanced progress bar with precision focus
|
1321 |
+
sep = np.mean(d_an) - np.mean(d_ap)
|
1322 |
+
actual_forgery_ratio = forgery_batch_count / (forgery_batch_count + genuine_batch_count) if (forgery_batch_count + genuine_batch_count) > 0 else 0
|
1323 |
+
|
1324 |
+
# Get current mining stats
|
1325 |
+
mining_pcts = metrics.get_mining_percentages()
|
1326 |
+
|
1327 |
+
pbar.set_postfix({
|
1328 |
+
"loss": f"{batch_loss.item():.3f}",
|
1329 |
+
"h_neg%": f"{mining_pcts.get('neg_mining_hard_pct', 0):.0f}",
|
1330 |
+
"h_pos%": f"{mining_pcts.get('pos_mining_hard_pct', 0):.0f}",
|
1331 |
+
"d_sep": f"{sep:.3f}",
|
1332 |
+
"FP": f"{batch_fp_count}",
|
1333 |
+
"FN": f"{batch_fn_count}",
|
1334 |
+
"margin": f"{final_margin:.3f}"
|
1335 |
+
})
|
1336 |
+
|
1337 |
+
# Periodic logging
|
1338 |
+
if self.global_step % self.config.log_frequency == 0:
|
1339 |
+
enhanced_stats = {
|
1340 |
+
**curriculum_stats,
|
1341 |
+
**mining_pcts,
|
1342 |
+
"actual_forgery_ratio": actual_forgery_ratio,
|
1343 |
+
"batch_false_positives": int(batch_fp_count),
|
1344 |
+
"batch_false_negatives": int(batch_fn_count),
|
1345 |
+
"final_margin": final_margin,
|
1346 |
+
"epoch_progress": epoch_progress
|
1347 |
+
}
|
1348 |
+
self._log_training_step(metrics, enhanced_stats, self.global_step)
|
1349 |
+
|
1350 |
+
# Memory cleanup
|
1351 |
+
del anchors, positives, negatives, a_emb, p_emb, n_emb
|
1352 |
+
torch.cuda.empty_cache()
|
1353 |
+
|
1354 |
+
# End-of-epoch mining summary
|
1355 |
+
mining_pcts = metrics.get_mining_percentages()
|
1356 |
+
print(f"\n[Epoch {epoch} Mining Summary]")
|
1357 |
+
print(f" Hard Negatives: {mining_pcts.get('neg_mining_hard_pct', 0):.1f}% | Semi: {mining_pcts.get('neg_mining_semi_pct', 0):.1f}% | Easy: {mining_pcts.get('neg_mining_easy_pct', 0):.1f}%")
|
1358 |
+
print(f" Hard Positives: {mining_pcts.get('pos_mining_hard_pct', 0):.1f}% | Semi: {mining_pcts.get('pos_mining_semi_pct', 0):.1f}% | Easy: {mining_pcts.get('pos_mining_easy_pct', 0):.1f}%")
|
1359 |
+
print(f" Overall Hard Ratios - Positives: {mining_pcts.get('hard_positive_ratio', 0):.1f}% | Negatives: {mining_pcts.get('hard_negative_ratio', 0):.1f}%")
|
1360 |
+
print(f" False Positive Rate: {mining_pcts.get('false_positive_rate', 0):.1f}% | False Negative Rate: {mining_pcts.get('false_negative_rate', 0):.1f}%")
|
1361 |
+
|
1362 |
+
if "adaptive_separation" in mining_pcts:
|
1363 |
+
print(f" Adaptive separation: {mining_pcts['adaptive_separation']:.3f} | Overlap: {mining_pcts['adaptive_overlap']:.3f}")
|
1364 |
+
|
1365 |
+
return metrics
|
1366 |
+
|
1367 |
+
def validate_epoch(self, model: nn.Module, val_loader: DataLoader,
|
1368 |
+
epoch: int) -> Tuple[float, float, Dict[str, float]]:
|
1369 |
+
"""Validate for one epoch."""
|
1370 |
+
model.eval()
|
1371 |
+
|
1372 |
+
val_distances = []
|
1373 |
+
val_labels = []
|
1374 |
+
val_embeddings = []
|
1375 |
+
val_person_ids = []
|
1376 |
+
val_loss_total = 0.0
|
1377 |
+
|
1378 |
+
with torch.no_grad():
|
1379 |
+
pbar = tqdm(val_loader, desc=f"[Val] Epoch {epoch}")
|
1380 |
+
|
1381 |
+
for img1, img2, labels, id1, id2 in pbar:
|
1382 |
+
# Move to device
|
1383 |
+
img1 = img1.to(self.device, non_blocking=True)
|
1384 |
+
img2 = img2.to(self.device, non_blocking=True)
|
1385 |
+
labels = labels.to(self.device, non_blocking=True)
|
1386 |
+
id1 = id1.to(self.device, non_blocking=True)
|
1387 |
+
id2 = id2.to(self.device, non_blocking=True)
|
1388 |
+
|
1389 |
+
# Forward pass
|
1390 |
+
emb1, emb2 = model(img1, img2)
|
1391 |
+
distances = F.pairwise_distance(emb1, emb2)
|
1392 |
+
|
1393 |
+
# Compute validation loss
|
1394 |
+
val_loss = self._compute_validation_loss(emb1, emb2, labels, id1, id2, model.criterion)
|
1395 |
+
val_loss_total += val_loss.item()
|
1396 |
+
|
1397 |
+
# Collect results
|
1398 |
+
val_distances.extend(distances.cpu().numpy())
|
1399 |
+
val_labels.extend(labels.cpu().numpy())
|
1400 |
+
val_embeddings.append(emb1.cpu().numpy())
|
1401 |
+
val_embeddings.append(emb2.cpu().numpy())
|
1402 |
+
val_person_ids.extend(id1.cpu().numpy())
|
1403 |
+
val_person_ids.extend(id2.cpu().numpy())
|
1404 |
+
|
1405 |
+
# Update progress
|
1406 |
+
pos_mask = labels == 1
|
1407 |
+
neg_mask = labels == 0
|
1408 |
+
pos_dist = distances[pos_mask].mean().item() if pos_mask.any() else 0.0
|
1409 |
+
neg_dist = distances[neg_mask].mean().item() if neg_mask.any() else 0.0
|
1410 |
+
|
1411 |
+
pbar.set_postfix({
|
1412 |
+
"loss": f"{val_loss.item():.4f}",
|
1413 |
+
"d_pos": f"{pos_dist:.3f}",
|
1414 |
+
"d_neg": f"{neg_dist:.3f}",
|
1415 |
+
"sep": f"{neg_dist - pos_dist:.3f}"
|
1416 |
+
})
|
1417 |
+
|
1418 |
+
# Memory cleanup
|
1419 |
+
del img1, img2, emb1, emb2
|
1420 |
+
torch.cuda.empty_cache()
|
1421 |
+
|
1422 |
+
# Process results
|
1423 |
+
val_distances = np.array(val_distances)
|
1424 |
+
val_labels = np.array(val_labels)
|
1425 |
+
val_embeddings = np.concatenate(val_embeddings)
|
1426 |
+
val_person_ids = np.array(val_person_ids)
|
1427 |
+
avg_val_loss = val_loss_total / len(val_loader)
|
1428 |
+
|
1429 |
+
# Compute metrics
|
1430 |
+
threshold, eer, metrics_dict = self._compute_validation_metrics(
|
1431 |
+
val_distances, val_labels, avg_val_loss
|
1432 |
+
)
|
1433 |
+
|
1434 |
+
# Update model threshold
|
1435 |
+
model.distance_threshold = threshold
|
1436 |
+
|
1437 |
+
return threshold, eer, {
|
1438 |
+
"metrics": metrics_dict,
|
1439 |
+
"embeddings": val_embeddings,
|
1440 |
+
"labels": np.repeat(val_labels, 2),
|
1441 |
+
"person_ids": val_person_ids,
|
1442 |
+
"distances": np.repeat(val_distances, 2)
|
1443 |
+
}
|
1444 |
+
|
1445 |
+
def _compute_validation_loss(self, emb1: torch.Tensor, emb2: torch.Tensor,
|
1446 |
+
binary_labels: torch.Tensor, person_ids1: torch.Tensor,
|
1447 |
+
person_ids2: torch.Tensor, criterion) -> torch.Tensor:
|
1448 |
+
"""Compute validation loss using MultiSimilarityLoss."""
|
1449 |
+
labels1 = person_ids1.clone()
|
1450 |
+
labels2 = person_ids2.clone()
|
1451 |
+
|
1452 |
+
# Handle forged pairs
|
1453 |
+
forged_mask = binary_labels == 0
|
1454 |
+
if forged_mask.any():
|
1455 |
+
max_person_id = torch.max(torch.cat([person_ids1, person_ids2])).item()
|
1456 |
+
labels2[forged_mask] = labels2[forged_mask] + max_person_id + 1
|
1457 |
+
|
1458 |
+
# Handle genuine pairs
|
1459 |
+
genuine_mask = binary_labels == 1
|
1460 |
+
labels2[genuine_mask] = labels1[genuine_mask]
|
1461 |
+
|
1462 |
+
# Combine embeddings and labels
|
1463 |
+
all_embeddings = torch.cat([emb1, emb2], dim=0)
|
1464 |
+
all_labels = torch.cat([labels1, labels2], dim=0)
|
1465 |
+
|
1466 |
+
return criterion(all_embeddings, all_labels)
|
1467 |
+
|
1468 |
+
def _compute_validation_metrics(self, distances: np.ndarray, labels: np.ndarray,
|
1469 |
+
val_loss: float) -> Tuple[float, float, Dict[str, float]]:
|
1470 |
+
"""Compute comprehensive validation metrics with precision focus."""
|
1471 |
+
# Compute EER and threshold
|
1472 |
+
similarity_scores = 1.0 / (distances + 1e-8)
|
1473 |
+
fpr, tpr, thresholds = roc_curve(labels, similarity_scores, pos_label=1)
|
1474 |
+
fnr = 1 - tpr
|
1475 |
+
eer_idx = np.nanargmin(np.abs(fpr - fnr))
|
1476 |
+
eer = fpr[eer_idx]
|
1477 |
+
eer_threshold = 1.0 / thresholds[eer_idx]
|
1478 |
+
|
1479 |
+
# Get precision-optimized threshold
|
1480 |
+
precision_threshold = self._compute_precision_optimized_threshold(distances, labels)
|
1481 |
+
|
1482 |
+
# Use precision-optimized threshold instead of EER threshold
|
1483 |
+
threshold = precision_threshold
|
1484 |
+
|
1485 |
+
# Compute metrics with precision-optimized threshold
|
1486 |
+
predictions = (distances < threshold).astype(int)
|
1487 |
+
precision, recall, f1, _ = precision_recall_fscore_support(
|
1488 |
+
labels, predictions, average='binary', zero_division=0
|
1489 |
+
)
|
1490 |
+
accuracy = (predictions == labels).mean()
|
1491 |
+
roc_auc = auc(fpr, tpr)
|
1492 |
+
|
1493 |
+
# Distance statistics
|
1494 |
+
genuine_dist = np.mean([d for d, l in zip(distances, labels) if l == 1])
|
1495 |
+
forged_dist = np.mean([d for d, l in zip(distances, labels) if l == 0])
|
1496 |
+
separation = forged_dist - genuine_dist
|
1497 |
+
|
1498 |
+
# Confidence scores
|
1499 |
+
confidences = 1.0 / (distances + 1e-8)
|
1500 |
+
conf_genuine = np.mean([c for c, l in zip(confidences, labels) if l == 1])
|
1501 |
+
conf_forged = np.mean([c for c, l in zip(confidences, labels) if l == 0])
|
1502 |
+
|
1503 |
+
metrics_dict = {
|
1504 |
+
"val_loss": val_loss,
|
1505 |
+
"val_EER": eer,
|
1506 |
+
"val_f1": f1,
|
1507 |
+
"val_auc": roc_auc,
|
1508 |
+
"val_accuracy": accuracy,
|
1509 |
+
"val_precision": precision,
|
1510 |
+
"val_recall": recall,
|
1511 |
+
"val_separation": separation,
|
1512 |
+
"val_genuine_dist": genuine_dist,
|
1513 |
+
"val_forged_dist": forged_dist,
|
1514 |
+
"val_genuine_conf": conf_genuine,
|
1515 |
+
"val_forged_conf": conf_forged,
|
1516 |
+
"threshold": threshold,
|
1517 |
+
"eer_threshold": eer_threshold,
|
1518 |
+
"precision_threshold": precision_threshold
|
1519 |
+
}
|
1520 |
+
|
1521 |
+
return threshold, eer, metrics_dict
|
1522 |
+
|
1523 |
+
def _log_training_step(self, metrics: TrainingMetrics, curriculum_stats: Dict, step: int):
|
1524 |
+
"""Log training step metrics."""
|
1525 |
+
if not metrics.losses:
|
1526 |
+
return
|
1527 |
+
|
1528 |
+
try:
|
1529 |
+
# Compute separation metrics
|
1530 |
+
sep_metrics = metrics.compute_separation_metrics()
|
1531 |
+
|
1532 |
+
# Get mining percentages
|
1533 |
+
mining_percentages = metrics.get_mining_percentages()
|
1534 |
+
|
1535 |
+
# Log to MLflow
|
1536 |
+
log_dict = {
|
1537 |
+
"train_loss": np.mean(metrics.losses[-10:]), # Last 10 batches
|
1538 |
+
**sep_metrics,
|
1539 |
+
**mining_percentages,
|
1540 |
+
**curriculum_stats,
|
1541 |
+
**metrics.learning_rates
|
1542 |
+
}
|
1543 |
+
|
1544 |
+
mlflow.log_metrics(log_dict, step=step)
|
1545 |
+
|
1546 |
+
except Exception as e:
|
1547 |
+
print(f"Warning: Failed to log training step metrics: {e}")
|
1548 |
+
|
1549 |
+
def _log_epoch_metrics(self, train_metrics: TrainingMetrics, val_metrics: Dict, epoch: int):
|
1550 |
+
"""Log comprehensive epoch metrics."""
|
1551 |
+
try:
|
1552 |
+
# Training metrics
|
1553 |
+
train_sep = train_metrics.compute_separation_metrics()
|
1554 |
+
train_mining = train_metrics.get_mining_percentages()
|
1555 |
+
|
1556 |
+
log_dict = {
|
1557 |
+
"epoch": epoch,
|
1558 |
+
"train_loss_epoch": np.mean(train_metrics.losses),
|
1559 |
+
**train_sep,
|
1560 |
+
**train_mining,
|
1561 |
+
**val_metrics["metrics"],
|
1562 |
+
**train_metrics.learning_rates
|
1563 |
+
}
|
1564 |
+
|
1565 |
+
mlflow.log_metrics(log_dict, step=epoch)
|
1566 |
+
|
1567 |
+
# Log key metrics to console
|
1568 |
+
self.logger.info(f"Epoch {epoch}/{self.config.max_epochs} Summary:")
|
1569 |
+
self.logger.info(f" Train Loss: {log_dict['train_loss_epoch']:.4f}")
|
1570 |
+
self.logger.info(f" Val EER: {log_dict['val_EER']:.4f}")
|
1571 |
+
self.logger.info(f" Val F1: {log_dict['val_f1']:.4f}")
|
1572 |
+
self.logger.info(f" Separation: {log_dict['separation']:.4f}")
|
1573 |
+
|
1574 |
+
except Exception as e:
|
1575 |
+
self.logger.error(f"Failed to log epoch metrics: {e}")
|
1576 |
+
# Log minimal metrics as fallback
|
1577 |
+
mlflow.log_metrics({
|
1578 |
+
"epoch": epoch,
|
1579 |
+
"train_loss_epoch": np.mean(train_metrics.losses) if train_metrics.losses else 0.0,
|
1580 |
+
**val_metrics["metrics"]
|
1581 |
+
}, step=epoch)
|
1582 |
+
|
1583 |
+
def _save_checkpoint(self, model: nn.Module, optimizer: torch.optim.Optimizer,
|
1584 |
+
scheduler: Any, scaler: torch.amp.GradScaler, epoch: int,
|
1585 |
+
threshold: float, eer: float, checkpoint_dir: str, is_best: bool = False):
|
1586 |
+
"""Save model checkpoint."""
|
1587 |
+
checkpoint = {
|
1588 |
+
"epoch": epoch,
|
1589 |
+
"model_state_dict": model.state_dict(),
|
1590 |
+
"optimizer_state_dict": optimizer.state_dict(),
|
1591 |
+
"scheduler_state_dict": scheduler.state_dict(),
|
1592 |
+
"scaler_state_dict": scaler.state_dict(),
|
1593 |
+
"prediction_threshold": threshold,
|
1594 |
+
"best_eer": self.best_eer,
|
1595 |
+
"eer": eer,
|
1596 |
+
"config": asdict(self.config)
|
1597 |
+
}
|
1598 |
+
|
1599 |
+
# Save regular checkpoint
|
1600 |
+
if epoch % self.config.save_frequency == 0:
|
1601 |
+
torch.save(checkpoint, os.path.join(checkpoint_dir, f"epoch_{epoch}.pth"))
|
1602 |
+
|
1603 |
+
# Save best checkpoint
|
1604 |
+
if is_best:
|
1605 |
+
torch.save(checkpoint, os.path.join(checkpoint_dir, "best_model.pth"))
|
1606 |
+
self.logger.info(f"New best model saved with EER: {eer:.4f}")
|
1607 |
+
|
1608 |
+
def _create_visualizations(self, val_results: Dict, epoch: int, figures_dir: str):
|
1609 |
+
"""Create comprehensive visualizations."""
|
1610 |
+
if epoch % self.config.visualize_frequency != 0:
|
1611 |
+
return
|
1612 |
+
|
1613 |
+
# Distance distribution plot
|
1614 |
+
self._plot_distance_distribution(
|
1615 |
+
val_results["distances"][:len(val_results["distances"])//2],
|
1616 |
+
val_results["labels"][:len(val_results["labels"])//2],
|
1617 |
+
epoch, figures_dir
|
1618 |
+
)
|
1619 |
+
|
1620 |
+
# t-SNE embedding visualization
|
1621 |
+
self._plot_tsne_embeddings(
|
1622 |
+
val_results["embeddings"],
|
1623 |
+
val_results["labels"],
|
1624 |
+
val_results["person_ids"],
|
1625 |
+
val_results["distances"],
|
1626 |
+
epoch, figures_dir
|
1627 |
+
)
|
1628 |
+
|
1629 |
+
def _plot_distance_distribution(self, distances: np.ndarray, labels: np.ndarray,
|
1630 |
+
epoch: int, figures_dir: str):
|
1631 |
+
"""Plot distance distribution."""
|
1632 |
+
genuine_dists = distances[labels == 1]
|
1633 |
+
forged_dists = distances[labels == 0]
|
1634 |
+
|
1635 |
+
plt.figure(figsize=(12, 8))
|
1636 |
+
plt.hist(genuine_dists, bins=50, alpha=0.6, color='blue',
|
1637 |
+
label=f'Genuine (μ={np.mean(genuine_dists):.4f}±{np.std(genuine_dists):.4f})')
|
1638 |
+
plt.hist(forged_dists, bins=50, alpha=0.6, color='red',
|
1639 |
+
label=f'Forged (μ={np.mean(forged_dists):.4f}±{np.std(forged_dists):.4f})')
|
1640 |
+
|
1641 |
+
separation = np.mean(forged_dists) - np.mean(genuine_dists)
|
1642 |
+
plt.axvline(np.mean(genuine_dists), color='blue', linestyle='--', alpha=0.7)
|
1643 |
+
plt.axvline(np.mean(forged_dists), color='red', linestyle='--', alpha=0.7)
|
1644 |
+
|
1645 |
+
plt.title(f'Distance Distribution - Epoch {epoch}\nSeparation: {separation:.4f}', fontsize=14)
|
1646 |
+
plt.xlabel('Embedding Distance', fontsize=12)
|
1647 |
+
plt.ylabel('Frequency', fontsize=12)
|
1648 |
+
plt.legend(fontsize=12)
|
1649 |
+
plt.grid(alpha=0.3)
|
1650 |
+
|
1651 |
+
plt.savefig(os.path.join(figures_dir, f"distance_dist_epoch_{epoch}.png"),
|
1652 |
+
dpi=150, bbox_inches='tight')
|
1653 |
+
plt.close()
|
1654 |
+
|
1655 |
+
|
1656 |
+
def _plot_tsne_embeddings(self, embeddings: np.ndarray, labels: np.ndarray,
|
1657 |
+
person_ids: np.ndarray, distances: np.ndarray,
|
1658 |
+
epoch: int, figures_dir: str, n_samples: int = 3000):
|
1659 |
+
"""Create comprehensive t-SNE visualization."""
|
1660 |
+
# Sample for computational efficiency
|
1661 |
+
if len(embeddings) > n_samples:
|
1662 |
+
indices = np.random.choice(len(embeddings), n_samples, replace=False)
|
1663 |
+
embeddings = embeddings[indices]
|
1664 |
+
labels = labels[indices]
|
1665 |
+
person_ids = person_ids[indices]
|
1666 |
+
distances = distances[indices]
|
1667 |
+
|
1668 |
+
# Compute t-SNE
|
1669 |
+
tsne = TSNE(n_components=2, random_state=42, perplexity=min(30, len(embeddings)-1))
|
1670 |
+
embeddings_2d = tsne.fit_transform(embeddings)
|
1671 |
+
|
1672 |
+
fig, axes = plt.subplots(1, 3, figsize=(20, 6))
|
1673 |
+
|
1674 |
+
# 1. Genuine vs Forged
|
1675 |
+
for label_val, color, name in [(0, 'red', 'Forged'), (1, 'blue', 'Genuine')]:
|
1676 |
+
mask = labels == label_val
|
1677 |
+
if mask.any():
|
1678 |
+
axes[0].scatter(embeddings_2d[mask, 0], embeddings_2d[mask, 1],
|
1679 |
+
c=color, label=name, alpha=0.6, s=20)
|
1680 |
+
axes[0].set_title(f'Genuine vs Forged - Epoch {epoch}')
|
1681 |
+
axes[0].legend()
|
1682 |
+
axes[0].grid(alpha=0.3)
|
1683 |
+
|
1684 |
+
# 2. Person clusters
|
1685 |
+
unique_ids = np.unique(person_ids)
|
1686 |
+
colors = plt.cm.tab20(np.linspace(0, 1, min(20, len(unique_ids))))
|
1687 |
+
|
1688 |
+
# Show top 15 most frequent IDs
|
1689 |
+
id_counts = {pid: np.sum(person_ids == pid) for pid in unique_ids}
|
1690 |
+
top_ids = sorted(id_counts.items(), key=lambda x: x[1], reverse=True)[:15]
|
1691 |
+
|
1692 |
+
for idx, (pid, count) in enumerate(top_ids):
|
1693 |
+
mask = person_ids == pid
|
1694 |
+
color = colors[idx % len(colors)]
|
1695 |
+
|
1696 |
+
# Plot the cluster points
|
1697 |
+
axes[1].scatter(embeddings_2d[mask, 0], embeddings_2d[mask, 1],
|
1698 |
+
c=[color], label=f'ID {pid} (n={count})', alpha=0.7, s=25)
|
1699 |
+
|
1700 |
+
# Compute the centroid (mean position) of the points in this cluster
|
1701 |
+
centroid = np.mean(embeddings_2d[mask], axis=0)
|
1702 |
+
|
1703 |
+
# Add the person ID text at the centroid
|
1704 |
+
axes[1].text(centroid[0], centroid[1], f'ID {pid}', fontsize=10, color='black', alpha=0.8, ha='center')
|
1705 |
+
|
1706 |
+
# Plot others in gray
|
1707 |
+
other_mask = ~np.isin(person_ids, [pid for pid, _ in top_ids])
|
1708 |
+
if other_mask.any():
|
1709 |
+
axes[1].scatter(embeddings_2d[other_mask, 0], embeddings_2d[other_mask, 1],
|
1710 |
+
c='gray', label='Others', alpha=0.3, s=15)
|
1711 |
+
|
1712 |
+
axes[1].set_title(f'Person Clusters - Epoch {epoch}')
|
1713 |
+
axes[1].legend(bbox_to_anchor=(1.05, 1), loc='upper left', fontsize=8)
|
1714 |
+
axes[1].grid(alpha=0.3)
|
1715 |
+
|
1716 |
+
# 3. Distance-based coloring
|
1717 |
+
scatter = axes[2].scatter(embeddings_2d[:, 0], embeddings_2d[:, 1],
|
1718 |
+
c=distances, cmap='viridis', alpha=0.7, s=20)
|
1719 |
+
plt.colorbar(scatter, ax=axes[2], label='Distance')
|
1720 |
+
axes[2].set_title(f'Distance Visualization - Epoch {epoch}')
|
1721 |
+
axes[2].grid(alpha=0.3)
|
1722 |
+
|
1723 |
+
plt.tight_layout()
|
1724 |
+
plt.savefig(os.path.join(figures_dir, f"tsne_epoch_{epoch}.png"),
|
1725 |
+
dpi=150, bbox_inches='tight')
|
1726 |
+
plt.close()
|
1727 |
+
|
1728 |
+
|
1729 |
+
def train(self):
|
1730 |
+
"""Main training loop."""
|
1731 |
+
torch.backends.cudnn.benchmark = True
|
1732 |
+
self.logger.info(f"Starting training on device: {self.device}")
|
1733 |
+
|
1734 |
+
# Prepare components
|
1735 |
+
train_dataset, val_dataset = self._prepare_datasets()
|
1736 |
+
model, optimizer, scheduler = self._setup_model_and_optimizer()
|
1737 |
+
scaler = torch.amp.GradScaler(self.device.type, enabled=(self.device.type == "cuda"))
|
1738 |
+
|
1739 |
+
# MLflow setup
|
1740 |
+
with self.mlflow_manager.start_run(run_id=self.config.run_id):
|
1741 |
+
run_id = mlflow.active_run().info.run_id
|
1742 |
+
self.mlflow_manager.log_config(self.config)
|
1743 |
+
|
1744 |
+
# Setup checkpoints
|
1745 |
+
checkpoint_dir, figures_dir = self._setup_checkpoint_management(run_id)
|
1746 |
+
|
1747 |
+
# Load checkpoint if specified
|
1748 |
+
start_epoch = self._load_checkpoint(model, optimizer, scheduler, scaler)
|
1749 |
+
|
1750 |
+
# Data loaders
|
1751 |
+
val_loader = DataLoader(
|
1752 |
+
val_dataset, batch_size=self.config.batch_size, shuffle=False,
|
1753 |
+
num_workers=4, pin_memory=True, prefetch_factor=2
|
1754 |
+
)
|
1755 |
+
|
1756 |
+
# Training loop
|
1757 |
+
for epoch in range(start_epoch, self.config.max_epochs + 1):
|
1758 |
+
self.current_epoch = epoch
|
1759 |
+
|
1760 |
+
# Update curriculum
|
1761 |
+
train_dataset.set_epoch(epoch)
|
1762 |
+
train_loader = DataLoader(
|
1763 |
+
train_dataset, batch_size=self.config.batch_size, shuffle=True,
|
1764 |
+
num_workers=4, pin_memory=True, persistent_workers=True, prefetch_factor=2
|
1765 |
+
)
|
1766 |
+
|
1767 |
+
# Training phase
|
1768 |
+
train_metrics = self.train_epoch(model, train_loader, optimizer, scaler, epoch)
|
1769 |
+
|
1770 |
+
# Validation phase
|
1771 |
+
threshold, eer, val_results = self.validate_epoch(model, val_loader, epoch)
|
1772 |
+
|
1773 |
+
# Logging
|
1774 |
+
self._log_epoch_metrics(train_metrics, val_results, epoch)
|
1775 |
+
|
1776 |
+
# Visualizations
|
1777 |
+
self._create_visualizations(val_results, epoch, figures_dir)
|
1778 |
+
|
1779 |
+
# Model checkpoint management
|
1780 |
+
is_best = eer < self.best_eer
|
1781 |
+
if is_best:
|
1782 |
+
self.best_eer = eer
|
1783 |
+
self.patience_counter = 0
|
1784 |
+
else:
|
1785 |
+
self.patience_counter += 1
|
1786 |
+
|
1787 |
+
self._save_checkpoint(
|
1788 |
+
model, optimizer, scheduler, scaler, epoch,
|
1789 |
+
threshold, eer, checkpoint_dir, is_best
|
1790 |
+
)
|
1791 |
+
|
1792 |
+
# Early stopping
|
1793 |
+
if self.patience_counter >= self.config.patience:
|
1794 |
+
self.logger.info(f"Early stopping after {self.config.patience} epochs without improvement")
|
1795 |
+
break
|
1796 |
+
|
1797 |
+
# Learning rate scheduling
|
1798 |
+
if self.config.lr_scheduler == "cosine":
|
1799 |
+
scheduler.step()
|
1800 |
+
else:
|
1801 |
+
scheduler.step(eer)
|
1802 |
+
|
1803 |
+
# Memory cleanup
|
1804 |
+
gc.collect()
|
1805 |
+
torch.cuda.empty_cache()
|
1806 |
+
|
1807 |
+
# Final logging
|
1808 |
+
mlflow.log_metric("final_best_eer", self.best_eer)
|
1809 |
+
self.logger.info(f"Training completed. Best EER: {self.best_eer:.4f}")
|
1810 |
+
|
1811 |
+
# ----------------------------
|
1812 |
+
# Image Processing Utilities
|
1813 |
+
# ----------------------------
|
1814 |
+
def estimate_background_color_pil(image: Image.Image, border_width: int = 10,
|
1815 |
+
method: str = "median") -> np.ndarray:
|
1816 |
+
"""Estimate background color from image borders."""
|
1817 |
+
if image.mode != 'RGB':
|
1818 |
+
image = image.convert('RGB')
|
1819 |
+
|
1820 |
+
np_img = np.array(image)
|
1821 |
+
h, w, _ = np_img.shape
|
1822 |
+
|
1823 |
+
# Extract border pixels
|
1824 |
+
top = np_img[:border_width, :, :].reshape(-1, 3)
|
1825 |
+
bottom = np_img[-border_width:, :, :].reshape(-1, 3)
|
1826 |
+
left = np_img[:, :border_width, :].reshape(-1, 3)
|
1827 |
+
right = np_img[:, -border_width:, :].reshape(-1, 3)
|
1828 |
+
|
1829 |
+
all_border_pixels = np.concatenate([top, bottom, left, right], axis=0)
|
1830 |
+
|
1831 |
+
if method == "mean":
|
1832 |
+
return np.mean(all_border_pixels, axis=0).astype(np.uint8)
|
1833 |
+
else:
|
1834 |
+
return np.median(all_border_pixels, axis=0).astype(np.uint8)
|
1835 |
+
|
1836 |
+
def replace_background_with_white(image_name: str, folder_img: str,
|
1837 |
+
tolerance: int = 40, method: str = "median",
|
1838 |
+
remove_bg: bool = False) -> Image.Image:
|
1839 |
+
"""Replace background with white based on border color estimation."""
|
1840 |
+
image_path = os.path.join(folder_img, image_name)
|
1841 |
+
image = Image.open(image_path).convert("RGB")
|
1842 |
+
|
1843 |
+
if not remove_bg:
|
1844 |
+
return image
|
1845 |
+
|
1846 |
+
np_img = np.array(image)
|
1847 |
+
bg_color = estimate_background_color_pil(image, method=method)
|
1848 |
+
|
1849 |
+
# Create mask for background pixels
|
1850 |
+
diff = np.abs(np_img.astype(np.int32) - bg_color.astype(np.int32))
|
1851 |
+
mask = np.all(diff < tolerance, axis=2)
|
1852 |
+
|
1853 |
+
# Replace background with white
|
1854 |
+
result = np_img.copy()
|
1855 |
+
result[mask] = [255, 255, 255]
|
1856 |
+
|
1857 |
+
return Image.fromarray(result)
|
1858 |
+
|
1859 |
+
# ----------------------------
|
1860 |
+
# Main Execution
|
1861 |
+
# ----------------------------
|
1862 |
+
def main():
|
1863 |
+
"""Main execution function with aggressive curriculum."""
|
1864 |
+
# Test distance ranges first
|
1865 |
+
print("\n[INFO] Testing distance ranges for margin calibration...")
|
1866 |
+
dummy_emb1 = F.normalize(torch.randn(1000, 128), p=2, dim=1)
|
1867 |
+
dummy_emb2 = F.normalize(torch.randn(1000, 128), p=2, dim=1)
|
1868 |
+
dummy_distances = F.pairwise_distance(dummy_emb1, dummy_emb2).numpy()
|
1869 |
+
print(f"Random embeddings: mean={dummy_distances.mean():.3f}, std={dummy_distances.std():.3f}")
|
1870 |
+
print(f"Expected margin range: {dummy_distances.std() * 0.5:.3f} - {dummy_distances.std() * 1.5:.3f}")
|
1871 |
+
|
1872 |
+
# Aggressive curriculum configuration
|
1873 |
+
CONFIG.model_name = "resnet34"
|
1874 |
+
CONFIG.embedding_dim = 128
|
1875 |
+
CONFIG.max_epochs = 20 # Shorter with aggressive curriculum
|
1876 |
+
CONFIG.head_lr = 2e-3 # Higher for faster adaptation
|
1877 |
+
CONFIG.backbone_lr = 1e-4
|
1878 |
+
CONFIG.curriculum_strategy = "progressive"
|
1879 |
+
|
1880 |
+
# AGGRESSIVE SETTINGS
|
1881 |
+
CONFIG.initial_hard_ratio = 0.4 # Start much higher
|
1882 |
+
CONFIG.final_hard_ratio = 0.85 # Target very high
|
1883 |
+
CONFIG.curriculum_warmup_epochs = 1 # Very short warmup
|
1884 |
+
CONFIG.batch_size = 256 # Smaller batches for more frequent updates
|
1885 |
+
CONFIG.grad_accum_steps = 8 # Smaller accumulation
|
1886 |
+
|
1887 |
+
CONFIG.tracking_uri = "http://127.0.0.1:5555"
|
1888 |
+
#CONFIG.run_id = "aa58e3a1f3314351bc1dd2b82ab156ad"
|
1889 |
+
#CONFIG.last_epoch_weights = "../../model/models_checkpoints/aa58e3a1f3314351bc1dd2b82ab156ad/best_model.pth"
|
1890 |
+
|
1891 |
+
trainer = SignatureTrainer(CONFIG)
|
1892 |
+
trainer.train()
|
1893 |
+
if __name__ == "__main__":
|
1894 |
+
main()
|
Training model - validation - organigram.zip
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:c635ee72dc87e1ac02c03a9aa720a4723c1d5130222f1fedb59758d277797764
|
3 |
+
size 251911552
|