|
|
from typing import Optional
|
|
|
|
|
|
import matplotlib.pyplot as plt
|
|
|
import numpy as np
|
|
|
|
|
|
try:
|
|
|
from torch.utils.tensorboard import SummaryWriter
|
|
|
except ModuleNotFoundError:
|
|
|
SummaryWriter = None
|
|
|
|
|
|
try:
|
|
|
import wandb
|
|
|
except ModuleNotFoundError:
|
|
|
wandb = None
|
|
|
|
|
|
plt.ioff()
|
|
|
|
|
|
PLOT_FILE_NAME = "metrics_plot.png"
|
|
|
|
|
|
|
|
|
def safe_index(arr, idx):
|
|
|
return arr[idx] if 0 <= idx < len(arr) else None
|
|
|
|
|
|
|
|
|
class MetricsPlotSink:
|
|
|
"""
|
|
|
The MetricsPlotSink class records training metrics and saves them to a plot.
|
|
|
|
|
|
Args:
|
|
|
output_dir (str): Directory where the plot will be saved.
|
|
|
"""
|
|
|
|
|
|
def __init__(self, output_dir: str):
|
|
|
self.output_dir = output_dir
|
|
|
self.history = []
|
|
|
|
|
|
def update(self, values: dict):
|
|
|
self.history.append(values)
|
|
|
|
|
|
def save(self):
|
|
|
if not self.history:
|
|
|
print("No data to plot.")
|
|
|
return
|
|
|
|
|
|
def get_array(key):
|
|
|
return np.array([h[key] for h in self.history if key in h])
|
|
|
|
|
|
epochs = get_array('epoch')
|
|
|
train_loss = get_array('train_loss')
|
|
|
test_loss = get_array('test_loss')
|
|
|
test_coco_eval = [h['test_coco_eval_bbox'] for h in self.history if 'test_coco_eval_bbox' in h]
|
|
|
ap50_90 = np.array([safe_index(x, 0) for x in test_coco_eval if x is not None], dtype=np.float32)
|
|
|
ap50 = np.array([safe_index(x, 1) for x in test_coco_eval if x is not None], dtype=np.float32)
|
|
|
ar50_90 = np.array([safe_index(x, 8) for x in test_coco_eval if x is not None], dtype=np.float32)
|
|
|
|
|
|
ema_coco_eval = [h['ema_test_coco_eval_bbox'] for h in self.history if 'ema_test_coco_eval_bbox' in h]
|
|
|
ema_ap50_90 = np.array([safe_index(x, 0) for x in ema_coco_eval if x is not None], dtype=np.float32)
|
|
|
ema_ap50 = np.array([safe_index(x, 1) for x in ema_coco_eval if x is not None], dtype=np.float32)
|
|
|
ema_ar50_90 = np.array([safe_index(x, 8) for x in ema_coco_eval if x is not None], dtype=np.float32)
|
|
|
|
|
|
fig, axes = plt.subplots(2, 2, figsize=(18, 12))
|
|
|
|
|
|
|
|
|
if len(epochs) > 0:
|
|
|
if len(train_loss):
|
|
|
axes[0][0].plot(epochs, train_loss, label='Training Loss', marker='o', linestyle='-')
|
|
|
if len(test_loss):
|
|
|
axes[0][0].plot(epochs, test_loss, label='Validation Loss', marker='o', linestyle='--')
|
|
|
axes[0][0].set_title('Training and Validation Loss')
|
|
|
axes[0][0].set_xlabel('Epoch Number')
|
|
|
axes[0][0].set_ylabel('Loss Value')
|
|
|
axes[0][0].legend()
|
|
|
axes[0][0].grid(True)
|
|
|
|
|
|
|
|
|
if ap50.size > 0 or ema_ap50.size > 0:
|
|
|
if ap50.size > 0:
|
|
|
axes[0][1].plot(epochs[:len(ap50)], ap50, marker='o', linestyle='-', label='Base Model')
|
|
|
if ema_ap50.size > 0:
|
|
|
axes[0][1].plot(epochs[:len(ema_ap50)], ema_ap50, marker='o', linestyle='--', label='EMA Model')
|
|
|
axes[0][1].set_title('Average Precision @0.50')
|
|
|
axes[0][1].set_xlabel('Epoch Number')
|
|
|
axes[0][1].set_ylabel('AP50')
|
|
|
axes[0][1].legend()
|
|
|
axes[0][1].grid(True)
|
|
|
|
|
|
|
|
|
if ap50_90.size > 0 or ema_ap50_90.size > 0:
|
|
|
if ap50_90.size > 0:
|
|
|
axes[1][0].plot(epochs[:len(ap50_90)], ap50_90, marker='o', linestyle='-', label='Base Model')
|
|
|
if ema_ap50_90.size > 0:
|
|
|
axes[1][0].plot(epochs[:len(ema_ap50_90)], ema_ap50_90, marker='o', linestyle='--', label='EMA Model')
|
|
|
axes[1][0].set_title('Average Precision @0.50:0.95')
|
|
|
axes[1][0].set_xlabel('Epoch Number')
|
|
|
axes[1][0].set_ylabel('AP')
|
|
|
axes[1][0].legend()
|
|
|
axes[1][0].grid(True)
|
|
|
|
|
|
|
|
|
if ar50_90.size > 0 or ema_ar50_90.size > 0:
|
|
|
if ar50_90.size > 0:
|
|
|
axes[1][1].plot(epochs[:len(ar50_90)], ar50_90, marker='o', linestyle='-', label='Base Model')
|
|
|
if ema_ar50_90.size > 0:
|
|
|
axes[1][1].plot(epochs[:len(ema_ar50_90)], ema_ar50_90, marker='o', linestyle='--', label='EMA Model')
|
|
|
axes[1][1].set_title('Average Recall @0.50:0.95')
|
|
|
axes[1][1].set_xlabel('Epoch Number')
|
|
|
axes[1][1].set_ylabel('AR')
|
|
|
axes[1][1].legend()
|
|
|
axes[1][1].grid(True)
|
|
|
|
|
|
plt.tight_layout()
|
|
|
plt.savefig(f"{self.output_dir}/{PLOT_FILE_NAME}")
|
|
|
plt.close(fig)
|
|
|
print(f"Results saved to {self.output_dir}/{PLOT_FILE_NAME}")
|
|
|
|
|
|
|
|
|
class MetricsTensorBoardSink:
|
|
|
"""
|
|
|
Training metrics via TensorBoard.
|
|
|
|
|
|
Args:
|
|
|
output_dir (str): Directory where TensorBoard logs will be written.
|
|
|
"""
|
|
|
|
|
|
def __init__(self, output_dir: str):
|
|
|
if SummaryWriter:
|
|
|
self.writer = SummaryWriter(log_dir=output_dir)
|
|
|
print(f"TensorBoard logging initialized. To monitor logs, use 'tensorboard --logdir {output_dir}' and open http://localhost:6006/ in browser.")
|
|
|
else:
|
|
|
self.writer = None
|
|
|
print("Unable to initialize TensorBoard. Logging is turned off for this session. Run 'pip install tensorboard' to enable logging.")
|
|
|
|
|
|
def update(self, values: dict):
|
|
|
if not self.writer:
|
|
|
return
|
|
|
|
|
|
epoch = values['epoch']
|
|
|
|
|
|
if 'train_loss' in values:
|
|
|
self.writer.add_scalar("Loss/Train", values['train_loss'], epoch)
|
|
|
if 'test_loss' in values:
|
|
|
self.writer.add_scalar("Loss/Test", values['test_loss'], epoch)
|
|
|
|
|
|
if 'test_coco_eval_bbox' in values:
|
|
|
coco_eval = values['test_coco_eval_bbox']
|
|
|
ap50_90 = safe_index(coco_eval, 0)
|
|
|
ap50 = safe_index(coco_eval, 1)
|
|
|
ar50_90 = safe_index(coco_eval, 8)
|
|
|
if ap50_90 is not None:
|
|
|
self.writer.add_scalar("Metrics/Base/AP50_90", ap50_90, epoch)
|
|
|
if ap50 is not None:
|
|
|
self.writer.add_scalar("Metrics/Base/AP50", ap50, epoch)
|
|
|
if ar50_90 is not None:
|
|
|
self.writer.add_scalar("Metrics/Base/AR50_90", ar50_90, epoch)
|
|
|
|
|
|
if 'ema_test_coco_eval_bbox' in values:
|
|
|
ema_coco_eval = values['ema_test_coco_eval_bbox']
|
|
|
ema_ap50_90 = safe_index(ema_coco_eval, 0)
|
|
|
ema_ap50 = safe_index(ema_coco_eval, 1)
|
|
|
ema_ar50_90 = safe_index(ema_coco_eval, 8)
|
|
|
if ema_ap50_90 is not None:
|
|
|
self.writer.add_scalar("Metrics/EMA/AP50_90", ema_ap50_90, epoch)
|
|
|
if ema_ap50 is not None:
|
|
|
self.writer.add_scalar("Metrics/EMA/AP50", ema_ap50, epoch)
|
|
|
if ema_ar50_90 is not None:
|
|
|
self.writer.add_scalar("Metrics/EMA/AR50_90", ema_ar50_90, epoch)
|
|
|
|
|
|
self.writer.flush()
|
|
|
|
|
|
def close(self):
|
|
|
if not self.writer:
|
|
|
return
|
|
|
|
|
|
self.writer.close()
|
|
|
|
|
|
class MetricsWandBSink:
|
|
|
"""
|
|
|
Training metrics via W&B.
|
|
|
|
|
|
Args:
|
|
|
output_dir (str): Directory where W&B logs will be written locally.
|
|
|
project (str, optional): Associate this training run with a W&B project. If None, W&B will generate a name based on the git repo name.
|
|
|
run (str, optional): W&B run name. If None, W&B will generate a random name.
|
|
|
config (dict, optional): Input parameters, like hyperparameters or data preprocessing settings for the run for later comparison.
|
|
|
"""
|
|
|
|
|
|
def __init__(self, output_dir: str, project: Optional[str] = None, run: Optional[str] = None, config: Optional[dict] = None):
|
|
|
self.output_dir = output_dir
|
|
|
if wandb:
|
|
|
self.run = wandb.init(
|
|
|
project=project,
|
|
|
name=run,
|
|
|
config=config,
|
|
|
dir=output_dir
|
|
|
)
|
|
|
print(f"W&B logging initialized. To monitor logs, open {wandb.run.url}.")
|
|
|
else:
|
|
|
self.run = None
|
|
|
print("Unable to initialize W&B. Logging is turned off for this session. Run 'pip install wandb' to enable logging.")
|
|
|
|
|
|
def update(self, values: dict):
|
|
|
if not wandb or not self.run:
|
|
|
return
|
|
|
|
|
|
epoch = values['epoch']
|
|
|
log_dict = {"epoch": epoch}
|
|
|
|
|
|
if 'train_loss' in values:
|
|
|
log_dict["Loss/Train"] = values['train_loss']
|
|
|
if 'test_loss' in values:
|
|
|
log_dict["Loss/Test"] = values['test_loss']
|
|
|
|
|
|
if 'test_coco_eval_bbox' in values:
|
|
|
coco_eval = values['test_coco_eval_bbox']
|
|
|
ap50_90 = safe_index(coco_eval, 0)
|
|
|
ap50 = safe_index(coco_eval, 1)
|
|
|
ar50_90 = safe_index(coco_eval, 8)
|
|
|
if ap50_90 is not None:
|
|
|
log_dict["Metrics/Base/AP50_90"] = ap50_90
|
|
|
if ap50 is not None:
|
|
|
log_dict["Metrics/Base/AP50"] = ap50
|
|
|
if ar50_90 is not None:
|
|
|
log_dict["Metrics/Base/AR50_90"] = ar50_90
|
|
|
|
|
|
if 'ema_test_coco_eval_bbox' in values:
|
|
|
ema_coco_eval = values['ema_test_coco_eval_bbox']
|
|
|
ema_ap50_90 = safe_index(ema_coco_eval, 0)
|
|
|
ema_ap50 = safe_index(ema_coco_eval, 1)
|
|
|
ema_ar50_90 = safe_index(ema_coco_eval, 8)
|
|
|
if ema_ap50_90 is not None:
|
|
|
log_dict["Metrics/EMA/AP50_90"] = ema_ap50_90
|
|
|
if ema_ap50 is not None:
|
|
|
log_dict["Metrics/EMA/AP50"] = ema_ap50
|
|
|
if ema_ar50_90 is not None:
|
|
|
log_dict["Metrics/EMA/AR50_90"] = ema_ar50_90
|
|
|
|
|
|
wandb.log(log_dict)
|
|
|
|
|
|
def close(self):
|
|
|
if not wandb or not self.run:
|
|
|
return
|
|
|
|
|
|
self.run.finish() |
