| import copy |
| import os |
| import random |
| import time |
|
|
| import pandas as pd |
| import plotly.express as px |
| from PIL import Image |
| import matplotlib.pyplot as plt |
| import torch |
| import torch.nn as nn |
| import torch.optim as optim |
| from tqdm import tqdm |
|
|
| from . import utils |
|
|
|
|
| def plot_eda( |
| subset_data, |
| image_dir="../coco_subset_images/images", |
| sample_size=9, |
| export_dir=None, |
| export_prefix="eda_", |
| save_grid_png=False, |
| ): |
| if not subset_data: |
| print("No data available for EDA.") |
| return None, None |
|
|
| counts = {} |
| for item in subset_data: |
| label = item.get("true_label") |
| counts[label] = counts.get(label, 0) + 1 |
|
|
| df = pd.DataFrame(counts.items(), columns=["Category", "Count"]).sort_values( |
| "Count", ascending=False |
| ) |
|
|
| fig = px.bar( |
| df, |
| x="Category", |
| y="Count", |
| title="Data Distribution per Category (COCO Subset)", |
| color="Category", |
| template="plotly_white", |
| ) |
| fig.show() |
|
|
| if export_dir: |
| os.makedirs(export_dir, exist_ok=True) |
| fig.write_json(os.path.join(export_dir, f"{export_prefix}dist.json")) |
|
|
| sample_size = min(sample_size, len(subset_data)) |
| samples = random.sample(subset_data, sample_size) |
|
|
| rows, cols = 3, 3 |
| fig_grid, axes = plt.subplots(rows, cols, figsize=(12, 12)) |
| axes = axes.flatten() |
|
|
| for idx, ax in enumerate(axes): |
| if idx >= len(samples): |
| ax.axis("off") |
| continue |
|
|
| sample = samples[idx] |
| image_id = sample.get("image_id") |
| img_path = os.path.join(image_dir, f"{image_id}.jpg") |
|
|
| try: |
| img = Image.open(img_path).convert("RGB") |
| ax.imshow(img) |
| ax.axis("off") |
|
|
| gt_index = sample.get("ground_truth_idx") |
| choices = sample.get("choices", []) |
| gt = choices[gt_index] if gt_index is not None else "" |
|
|
| distractors = [choice for j, choice in enumerate(choices) if j != gt_index] |
| distractor = distractors[0] if distractors else "No distractor" |
|
|
| caption = f"GT: {gt[:30]}...\nDist: {distractor[:30]}..." |
| ax.set_title(caption, fontsize=9) |
| except Exception: |
| ax.set_title("Image missing", color="crimson", fontsize=10) |
| ax.axis("off") |
|
|
| fig_grid.suptitle("Random 3x3 Image Grid with Captions", fontsize=14) |
| fig_grid.tight_layout(rect=[0, 0, 1, 0.95]) |
| if save_grid_png and export_dir: |
| fig_grid.savefig(os.path.join(export_dir, f"{export_prefix}grid.png"), dpi=150) |
| plt.show() |
|
|
| return fig, fig_grid |
|
|
|
|
| def train_linear_probe( |
| model, |
| model_name, |
| preprocess, |
| k, |
| train_max_data, |
| eval_data, |
| config, |
| device, |
| data_loader, |
| num_epochs=10, |
| train_mode=True, |
| image_dir=None, |
| models_dir="../models", |
| wandb=None, |
| prev_k=None, |
| seed=42, |
| only_new=False, |
| warm_start_from_prev=False, |
| warm_start_lr_scale=1.0, |
| ): |
| if image_dir is None: |
| image_dir = "../coco_subset_images/images" |
|
|
| training_cfg = config.get("training", {}) if config else {} |
| label_smoothing = training_cfg.get("label_smoothing", 0.2) |
| dropout_p = training_cfg.get("dropout_p", 0.6) |
| early_stopping_patience = training_cfg.get("early_stopping_patience", 3) |
| min_delta = training_cfg.get("min_delta", 0.0) |
| bottleneck_ratio = getattr(model, "bottleneck_ratio", 32) |
| k_train_loader, k_test_loader = data_loader.create_dataloaders( |
| train_max_data, |
| eval_data, |
| image_dir, |
| preprocess, |
| batch_size=training_cfg.get("batch_size", 32), |
| k=k, |
| prev_k=prev_k, |
| seed=seed, |
| only_new=only_new, |
| ) |
|
|
| hidden_dim = max(1, model.dim // bottleneck_ratio) |
| model.img_gate = nn.Sequential( |
| nn.Linear(model.dim, hidden_dim, bias=False), |
| nn.ReLU(inplace=True), |
| nn.Linear(hidden_dim, model.dim, bias=False), |
| ).to(device) |
|
|
| model.txt_gate = nn.Sequential( |
| nn.Linear(model.dim, hidden_dim, bias=False), |
| nn.ReLU(inplace=True), |
| nn.Linear(hidden_dim, model.dim, bias=False), |
| ).to(device) |
|
|
| model.dropout = nn.Dropout(p=dropout_p).to(device) |
|
|
| if not os.path.exists(models_dir): |
| os.makedirs(models_dir) |
| save_path = os.path.join(models_dir, f"best_{model_name}_{k}shot.pth") |
|
|
| best_acc = 0.0 |
| start_epoch = 0 |
|
|
| warm_started = False |
| if os.path.exists(save_path): |
| print(f"Loading existing checkpoint for {model_name} from {save_path}...") |
| checkpoint = torch.load(save_path, map_location=device) |
| if "img_gate" in checkpoint: |
| try: |
| model.img_gate.load_state_dict(checkpoint["img_gate"]) |
| model.txt_gate.load_state_dict(checkpoint["txt_gate"]) |
| best_acc = checkpoint.get("best_acc", 0.0) |
| start_epoch = checkpoint.get("epoch", -1) + 1 |
| except RuntimeError: |
| print("Checkpoint incompatible (shape mismatch), reinitializing gates") |
| nn.init.zeros_(model.img_gate[2].weight) |
| nn.init.zeros_(model.txt_gate[2].weight) |
| else: |
| print("Checkpoint incompatible (legacy format), using initial zeros") |
| nn.init.zeros_(model.img_gate[2].weight) |
| nn.init.zeros_(model.txt_gate[2].weight) |
| elif warm_start_from_prev and prev_k is not None: |
| prev_path = os.path.join(models_dir, f"best_{model_name}_{prev_k}shot.pth") |
| if os.path.exists(prev_path): |
| print(f"Warm-starting {model_name} from {prev_path}...") |
| checkpoint = torch.load(prev_path, map_location=device) |
| if "img_gate" in checkpoint: |
| try: |
| model.img_gate.load_state_dict(checkpoint["img_gate"]) |
| model.txt_gate.load_state_dict(checkpoint["txt_gate"]) |
| warm_started = True |
| except RuntimeError: |
| print("Checkpoint incompatible (shape mismatch), using initial zeros") |
| nn.init.zeros_(model.img_gate[2].weight) |
| nn.init.zeros_(model.txt_gate[2].weight) |
| else: |
| print("Checkpoint incompatible (legacy format), using initial zeros") |
| nn.init.zeros_(model.img_gate[2].weight) |
| nn.init.zeros_(model.txt_gate[2].weight) |
| else: |
| if not train_mode: |
| print( |
| f"Warning: No saved checkpoint found at {save_path}. Random Zero-shot inferences will be served." |
| ) |
| else: |
| print(f"No checkpoint found. Initializing new Residual Gates for {model_name}.") |
| nn.init.zeros_(model.img_gate[2].weight) |
| nn.init.zeros_(model.txt_gate[2].weight) |
|
|
| if not train_mode: |
| print(f"--- Inferencing Only for {model_name} k={k} ---") |
| eval_acc, f1_macro, f1_micro, f1_weighted = utils.evaluate_model( |
| model, k_test_loader, device |
| ) |
| print( |
| "Loaded Checkpoint - Eval Acc: {:.4f} - F1 Macro: {:.4f} - F1 Micro: {:.4f} - F1 Weighted: {:.4f}".format( |
| eval_acc, f1_macro, f1_micro, f1_weighted |
| ) |
| ) |
| return eval_acc, f1_macro, f1_micro, f1_weighted |
|
|
| base_lr = training_cfg.get("learning_rate", 1e-3) |
| if warm_started: |
| base_lr = base_lr * warm_start_lr_scale |
|
|
| optimizer = optim.Adam( |
| list(model.img_gate.parameters()) + list(model.txt_gate.parameters()), |
| lr=base_lr, |
| weight_decay=training_cfg.get("weight_decay", 1e-3), |
| ) |
| scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs) |
| criterion = nn.CrossEntropyLoss(label_smoothing=label_smoothing) |
| best_model_state = None |
|
|
| patience = early_stopping_patience |
| epochs_no_improve = 0 |
|
|
| for epoch in range(start_epoch, num_epochs): |
| model.train() |
| total_loss = 0 |
| correct = 0 |
| total = 0 |
|
|
| pbar = tqdm(k_train_loader, desc=f"Epoch {epoch+1}/{num_epochs}") |
| for batch in pbar: |
| images, text_candidates, targets, _, _ = batch |
| images = images.to(device) |
| targets = targets.to(device) |
|
|
| optimizer.zero_grad() |
| logits = model(images, text_candidates) |
| loss = criterion(logits, targets) |
|
|
| loss.backward() |
| optimizer.step() |
|
|
| total_loss += loss.item() |
| preds = torch.argmax(logits, dim=1) |
| correct += (preds == targets).sum().item() |
| total += targets.size(0) |
| pbar.set_postfix({"loss": loss.item()}) |
|
|
| scheduler.step() |
|
|
| train_acc = correct / max(total, 1) |
| eval_acc, f1_macro, f1_micro, f1_weighted = utils.evaluate_model( |
| model, k_test_loader, device |
| ) |
|
|
| print( |
| "Epoch {} - Loss: {:.4f} - Train Acc: {:.4f} - Eval Acc: {:.4f} - F1 Macro: {:.4f} - F1 Micro: {:.4f} - F1 Weighted: {:.4f}".format( |
| epoch + 1, |
| total_loss / len(k_train_loader), |
| train_acc, |
| eval_acc, |
| f1_macro, |
| f1_micro, |
| f1_weighted, |
| ) |
| ) |
|
|
| if wandb is not None and wandb.run is not None: |
| wandb.log( |
| { |
| f"{model_name}_{k}shot_train_loss": total_loss / len(k_train_loader), |
| f"{model_name}_{k}shot_train_acc": train_acc, |
| f"{model_name}_{k}shot_eval_acc": eval_acc, |
| f"{model_name}_{k}shot_eval_f1_macro": f1_macro, |
| f"{model_name}_{k}shot_eval_f1_micro": f1_micro, |
| f"{model_name}_{k}shot_eval_f1_weighted": f1_weighted, |
| f"{model_name}_{k}shot_lr": scheduler.get_last_lr()[0], |
| } |
| ) |
|
|
| if eval_acc > best_acc + min_delta: |
| best_acc = eval_acc |
| best_model_state = { |
| "epoch": epoch, |
| "best_acc": best_acc, |
| "img_gate": copy.deepcopy(model.img_gate.state_dict()), |
| "txt_gate": copy.deepcopy(model.txt_gate.state_dict()), |
| } |
| torch.save(best_model_state, save_path) |
| epochs_no_improve = 0 |
| else: |
| epochs_no_improve += 1 |
|
|
| if epochs_no_improve >= patience: |
| print( |
| "Early stopping triggered at epoch {}. No improvement >= {:.4f} for {} epochs.".format( |
| epoch + 1, |
| min_delta, |
| patience, |
| ) |
| ) |
| break |
|
|
| if best_model_state: |
| model.img_gate.load_state_dict(best_model_state["img_gate"]) |
| model.txt_gate.load_state_dict(best_model_state["txt_gate"]) |
|
|
| if wandb is not None and wandb.run is not None: |
| artifact = wandb.Artifact(f"{model_name}_{k}shot_model", type="model") |
| artifact.add_file(save_path) |
| wandb.log_artifact(artifact) |
|
|
| return best_acc, f1_macro, f1_micro, f1_weighted |
|
|
|
|
| def run_few_shot_training( |
| vit_lp_model, |
| rn50_lp_model, |
| preprocess_vit, |
| preprocess_rn50, |
| train_max_data, |
| eval_data, |
| config, |
| device, |
| data_loader, |
| k_shots_list=None, |
| train_mode=True, |
| image_dir=None, |
| models_dir="../models", |
| wandb=None, |
| warm_start=False, |
| only_new=False, |
| seed=42, |
| warm_start_lr_scale=1.0, |
| ): |
| if k_shots_list is None: |
| k_shots_list = config.get("training", {}).get("k_shots", [8, 16, 32]) |
|
|
| if train_mode and wandb is not None: |
| try: |
| wandb.login() |
| project_name = ( |
| config.get("wandb", {}).get("project") |
| if config and "wandb" in config |
| else "coco_multimodal" |
| ) |
| wandb.init(project=project_name, entity=config.get("wandb", {}).get("entity")) |
| except Exception as exc: |
| print(f"\n[Warning] Weights & Biases failed to initialize: {exc}") |
|
|
| k_shots_list = sorted(k_shots_list) |
| prev_k = None |
| for k in k_shots_list: |
| print(f"\n--- Model Processor for ViT-B/32 k={k} ---") |
| train_linear_probe( |
| vit_lp_model, |
| "vit_b32", |
| preprocess_vit, |
| k, |
| train_max_data, |
| eval_data, |
| config, |
| device, |
| data_loader, |
| num_epochs=config.get("training", {}).get("epochs", 25), |
| train_mode=train_mode, |
| image_dir=image_dir, |
| models_dir=models_dir, |
| wandb=wandb, |
| prev_k=prev_k, |
| seed=seed, |
| only_new=only_new, |
| warm_start_from_prev=warm_start, |
| warm_start_lr_scale=warm_start_lr_scale, |
| ) |
|
|
| print(f"\n--- Model Processor for RN50 k={k} ---") |
| train_linear_probe( |
| rn50_lp_model, |
| "rn50", |
| preprocess_rn50, |
| k, |
| train_max_data, |
| eval_data, |
| config, |
| device, |
| data_loader, |
| num_epochs=config.get("training", {}).get("epochs", 25), |
| train_mode=train_mode, |
| image_dir=image_dir, |
| models_dir=models_dir, |
| wandb=wandb, |
| prev_k=prev_k, |
| seed=seed, |
| only_new=only_new, |
| warm_start_from_prev=warm_start, |
| warm_start_lr_scale=warm_start_lr_scale, |
| ) |
|
|
| prev_k = k |
|
|
| if train_mode and wandb is not None and wandb.run is not None: |
| wandb.finish() |
|
|
|
|
| def evaluate_zero_shot( |
| model, |
| preprocess, |
| train_max_data, |
| eval_data, |
| config, |
| device, |
| data_loader, |
| image_dir=None, |
| ): |
| if not eval_data: |
| return 0, 0, 0 |
|
|
| if image_dir is None: |
| image_dir = "../coco_subset_images/images" |
|
|
| _, test_loader = data_loader.create_dataloaders( |
| train_max_data, |
| eval_data, |
| image_dir, |
| preprocess, |
| batch_size=config.get("training", {}).get("batch_size", 32), |
| k=8, |
| ) |
|
|
| start_time = time.time() |
| acc, f1_macro, f1_micro, f1_weighted = utils.evaluate_model( |
| model, test_loader, device |
| ) |
| inference_time = time.time() - start_time |
|
|
| eval_count = len(eval_data) |
| print(f"Zero-Shot Accuracy: {acc:.4f}") |
| print(f"Zero-Shot F1 Macro: {f1_macro:.4f}") |
| print(f"Zero-Shot F1 Micro: {f1_micro:.4f}") |
| print(f"Zero-Shot F1 Weighted: {f1_weighted:.4f}") |
| print(f"Inference Time for {eval_count} samples: {inference_time:.4f}s") |
|
|
| return acc, f1_macro, f1_micro, f1_weighted, inference_time |
|
|
|
|
| def _ensure_probe_gates(model, device): |
| if not hasattr(model, "img_gate") or not hasattr(model, "txt_gate"): |
| bottleneck_ratio = getattr(model, "bottleneck_ratio", 32) |
| hidden_dim = max(1, model.dim // bottleneck_ratio) |
| model.img_gate = nn.Sequential( |
| nn.Linear(model.dim, hidden_dim, bias=False), |
| nn.ReLU(inplace=True), |
| nn.Linear(hidden_dim, model.dim, bias=False), |
| ).to(device) |
|
|
| model.txt_gate = nn.Sequential( |
| nn.Linear(model.dim, hidden_dim, bias=False), |
| nn.ReLU(inplace=True), |
| nn.Linear(hidden_dim, model.dim, bias=False), |
| ).to(device) |
|
|
| nn.init.zeros_(model.img_gate[2].weight) |
| nn.init.zeros_(model.txt_gate[2].weight) |
| else: |
| model.img_gate = model.img_gate.to(device) |
| model.txt_gate = model.txt_gate.to(device) |
|
|
|
|
| def _load_probe_checkpoint(model, model_name, k, device, models_dir): |
| save_path = os.path.join(models_dir, f"best_{model_name}_{k}shot.pth") |
| if not os.path.exists(save_path): |
| print(f"Checkpoint not found: {save_path}") |
| return False |
|
|
| checkpoint = torch.load(save_path, map_location=device) |
| if "img_gate" not in checkpoint or "txt_gate" not in checkpoint: |
| print(f"Incompatible checkpoint format: {save_path}") |
| return False |
|
|
| _ensure_probe_gates(model, device) |
| try: |
| model.img_gate.load_state_dict(checkpoint["img_gate"]) |
| model.txt_gate.load_state_dict(checkpoint["txt_gate"]) |
| except RuntimeError: |
| print(f"Checkpoint incompatible (shape mismatch): {save_path}") |
| return False |
| return True |
|
|
|
|
| def evaluate_few_shot( |
| model, |
| model_name, |
| preprocess, |
| k, |
| train_max_data, |
| eval_data, |
| config, |
| device, |
| data_loader, |
| image_dir=None, |
| models_dir="../models", |
| ): |
| if image_dir is None: |
| image_dir = "../coco_subset_images/images" |
|
|
| _, test_loader = data_loader.create_dataloaders( |
| train_max_data, |
| eval_data, |
| image_dir, |
| preprocess, |
| batch_size=config.get("training", {}).get("batch_size", 32), |
| k=k, |
| ) |
|
|
| if not _load_probe_checkpoint(model, model_name, k, device, models_dir): |
| return None |
|
|
| start_time = time.time() |
| acc, f1_macro, f1_micro, f1_weighted = utils.evaluate_model( |
| model, test_loader, device |
| ) |
| inference_time = time.time() - start_time |
|
|
| eval_count = len(eval_data) |
| print(f"{model_name.upper()} {k}-Shot Accuracy: {acc:.4f}") |
| print(f"{model_name.upper()} {k}-Shot F1 Macro: {f1_macro:.4f}") |
| print(f"{model_name.upper()} {k}-Shot F1 Micro: {f1_micro:.4f}") |
| print(f"{model_name.upper()} {k}-Shot F1 Weighted: {f1_weighted:.4f}") |
| print(f"Inference Time for {eval_count} samples: {inference_time:.4f}s") |
|
|
| return acc, f1_macro, f1_micro, f1_weighted, inference_time |
|
|
|
|
| def run_evaluation( |
| vit_zs_model, |
| rn50_zs_model, |
| vit_lp_model, |
| rn50_lp_model, |
| preprocess_vit, |
| preprocess_rn50, |
| train_max_data, |
| eval_data, |
| config, |
| device, |
| data_loader, |
| k_shots_list=None, |
| image_dir=None, |
| models_dir="../models", |
| export_dir=None, |
| export_prefix="", |
| ): |
| if k_shots_list is None: |
| k_shots_list = config.get("training", {}).get("k_shots", [8, 16, 32]) |
|
|
| results = [] |
|
|
| print("\n--- ViT-B/32 Zero-Shot Evaluation ---") |
| vit_zs_acc, vit_zs_f1_macro, vit_zs_f1_micro, vit_zs_f1_weighted, vit_zs_time = evaluate_zero_shot( |
| vit_zs_model, |
| preprocess_vit, |
| train_max_data, |
| eval_data, |
| config, |
| device, |
| data_loader, |
| image_dir=image_dir, |
| ) |
| results.append( |
| { |
| "Model": "ViT-B/32", |
| "Shots": 0, |
| "Mode": "Zero-shot", |
| "Accuracy": vit_zs_acc, |
| "F1Macro": vit_zs_f1_macro, |
| "F1Micro": vit_zs_f1_micro, |
| "F1Weighted": vit_zs_f1_weighted, |
| "TimeSec": vit_zs_time, |
| } |
| ) |
|
|
| print("\n--- RN50 Zero-Shot Evaluation ---") |
| rn50_zs_acc, rn50_zs_f1_macro, rn50_zs_f1_micro, rn50_zs_f1_weighted, rn50_zs_time = evaluate_zero_shot( |
| rn50_zs_model, |
| preprocess_rn50, |
| train_max_data, |
| eval_data, |
| config, |
| device, |
| data_loader, |
| image_dir=image_dir, |
| ) |
| results.append( |
| { |
| "Model": "RN50", |
| "Shots": 0, |
| "Mode": "Zero-shot", |
| "Accuracy": rn50_zs_acc, |
| "F1Macro": rn50_zs_f1_macro, |
| "F1Micro": rn50_zs_f1_micro, |
| "F1Weighted": rn50_zs_f1_weighted, |
| "TimeSec": rn50_zs_time, |
| } |
| ) |
|
|
| for k in k_shots_list: |
| print(f"\n--- ViT-B/32 Few-Shot Evaluation (k={k}) ---") |
| vit_res = evaluate_few_shot( |
| vit_lp_model, |
| "vit_b32", |
| preprocess_vit, |
| k, |
| train_max_data, |
| eval_data, |
| config, |
| device, |
| data_loader, |
| image_dir=image_dir, |
| models_dir=models_dir, |
| ) |
| if vit_res: |
| acc, f1_macro, f1_micro, f1_weighted, t = vit_res |
| results.append( |
| { |
| "Model": "ViT-B/32", |
| "Shots": k, |
| "Mode": "Few-shot", |
| "Accuracy": acc, |
| "F1Macro": f1_macro, |
| "F1Micro": f1_micro, |
| "F1Weighted": f1_weighted, |
| "TimeSec": t, |
| } |
| ) |
|
|
| print(f"\n--- RN50 Few-Shot Evaluation (k={k}) ---") |
| rn50_res = evaluate_few_shot( |
| rn50_lp_model, |
| "rn50", |
| preprocess_rn50, |
| k, |
| train_max_data, |
| eval_data, |
| config, |
| device, |
| data_loader, |
| image_dir=image_dir, |
| models_dir=models_dir, |
| ) |
| if rn50_res: |
| acc, f1_macro, f1_micro, f1_weighted, t = rn50_res |
| results.append( |
| { |
| "Model": "RN50", |
| "Shots": k, |
| "Mode": "Few-shot", |
| "Accuracy": acc, |
| "F1Macro": f1_macro, |
| "F1Micro": f1_micro, |
| "F1Weighted": f1_weighted, |
| "TimeSec": t, |
| } |
| ) |
|
|
| results_df = pd.DataFrame(results) |
| if results_df.empty: |
| print("No evaluation results to plot.") |
| return results_df |
|
|
| results_df["Shots"] = results_df["Shots"].astype(int) |
| eval_count = max(len(eval_data), 1) |
| results_df["TimePerSampleMs"] = (results_df["TimeSec"] / eval_count) * 1000.0 |
| results_df = results_df.sort_values(["Model", "Shots"]).reset_index(drop=True) |
|
|
| fig_acc = px.line( |
| results_df, |
| x="Shots", |
| y="Accuracy", |
| color="Model", |
| markers=True, |
| hover_data=["Mode"], |
| title="Accuracy vs Shots", |
| template="plotly_white", |
| ) |
| fig_acc.show() |
|
|
| fig_f1 = px.line( |
| results_df, |
| x="Shots", |
| y="F1Macro", |
| color="Model", |
| markers=True, |
| hover_data=["Mode"], |
| title="Macro F1 vs Shots", |
| template="plotly_white", |
| ) |
| fig_f1.show() |
|
|
| fig_f1_micro = px.line( |
| results_df, |
| x="Shots", |
| y="F1Micro", |
| color="Model", |
| markers=True, |
| hover_data=["Mode"], |
| title="Micro F1 vs Shots", |
| template="plotly_white", |
| ) |
| fig_f1_micro.show() |
|
|
| fig_f1_weighted = px.line( |
| results_df, |
| x="Shots", |
| y="F1Weighted", |
| color="Model", |
| markers=True, |
| hover_data=["Mode"], |
| title="Weighted F1 vs Shots", |
| template="plotly_white", |
| ) |
| fig_f1_weighted.show() |
|
|
| fig_time = px.bar( |
| results_df, |
| x="Shots", |
| y="TimePerSampleMs", |
| color="Model", |
| barmode="group", |
| hover_data=["Mode", "TimeSec"], |
| title="Inference Time per Sample (ms)", |
| template="plotly_white", |
| ) |
| fig_time.show() |
|
|
| if export_dir: |
| os.makedirs(export_dir, exist_ok=True) |
| fig_acc.write_json(os.path.join(export_dir, f"{export_prefix}acc.json")) |
| fig_f1.write_json(os.path.join(export_dir, f"{export_prefix}f1_macro.json")) |
| fig_f1_micro.write_json(os.path.join(export_dir, f"{export_prefix}f1_micro.json")) |
| fig_f1_weighted.write_json(os.path.join(export_dir, f"{export_prefix}f1_weighted.json")) |
| fig_time.write_json(os.path.join(export_dir, f"{export_prefix}time.json")) |
|
|
| return results_df |
|
|
