Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -169,3 +169,14 @@ cython_debug/
# and can be added to the global gitignore or merged into this file. For a more nuclear
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
.idea/

jobs_testing/
jobs/
.claude/

# Local outputs and external frameworks
images/
result_images/
thesis_figures/
trash/
funnybirds-framework/
44 changes: 35 additions & 9 deletions ciao/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,14 +75,25 @@ def _log_trajectory(run_id: str, results: ExplanationResult) -> None:
"""Batch-log per-region trajectory points to MLflow."""
ts_ms = int(time.time() * 1000)
trajectory_metrics = [
Metric(
key=f"region_{idx}/trajectory_best_score",
value=item["best_score"],
timestamp=ts_ms,
step=int(item["evals"]),
)
metric
for idx, region in enumerate(results.regions)
for item in region.trajectory
# Deduplicate by evals: multiple iterations may share the same eval_count
# (e.g. cached terminal hits in MCGS), which would violate MLflow's metric PK.
for item in {item["evals"]: item for item in region.trajectory}.values()
for metric in (
Metric(
key=f"region_{idx}/trajectory_best_score",
value=item["best_score"],
timestamp=ts_ms,
step=int(item["evals"]),
),
Metric(
key=f"region_{idx}/trajectory_time",
value=item["time"],
timestamp=ts_ms,
step=int(item["evals"]),
),
)
]
if not trajectory_metrics:
return
Expand All @@ -96,15 +107,18 @@ def _log_trajectory(run_id: str, results: ExplanationResult) -> None:


def _log_explanation_results(
run_id: str, results: ExplanationResult, elapsed: float
run_id: str,
results: ExplanationResult,
elapsed: float,
) -> None:
"""Log explanation params, per-region metrics, trajectory, and timing to MLflow."""
"""Log explanation params, baseline log-odds, per-region metrics, trajectory, and timing to MLflow."""
mlflow.log_params(
{
"target_class_idx": results.target_class_idx,
"class_name": results.class_name,
}
)
mlflow.log_metric("original_log_odds", results.original_log_odds)

for idx, region in enumerate(results.regions):
mlflow.log_metrics(
Expand All @@ -128,6 +142,17 @@ def _log_explanation_results(
f"region_{idx}/segments.json",
)

if (
results.combined_score is not None
and results.combined_probability_drop is not None
):
mlflow.log_metrics(
{
"all_regions/final_score": results.combined_score,
"all_regions/probability_drop": results.combined_probability_drop,
}
)

_log_trajectory(run_id, results)

mlflow.log_metric("time_seconds", elapsed)
Expand Down Expand Up @@ -207,6 +232,7 @@ def main(cfg: DictConfig) -> None:
predictor=predictor,
target_class_idx=cfg.target_class_idx,
max_regions=cfg.max_regions,
sigma=cfg.sigma,
desired_length=cfg.desired_length,
batch_size=cfg.batch_size,
segmentation=segmentation,
Expand Down
45 changes: 40 additions & 5 deletions ciao/algorithm/builder.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
"""Unified region builder orchestrating different search algorithms."""

from typing import Literal

import torch

from ciao.algorithm.context import SearchContext
Expand All @@ -9,15 +11,44 @@
from ciao.typing import ExplanationMethodFn


SeedSelectionMode = Literal[-1, 1] | None


def _select_seed_and_sign(
scores: dict[int, float], available_segments: list[int], sigma: SeedSelectionMode
) -> tuple[int, int]:
if not available_segments:
raise ValueError(
"Cannot select a seed from an empty set of available segments."
)

if sigma is None:
seed_idx = max(available_segments, key=lambda x: abs(scores[x]))
optimization_sign = 1 if scores[seed_idx] >= 0 else -1
elif sigma == 1:
seed_idx = max(available_segments, key=lambda x: scores[x])
optimization_sign = 1
elif sigma == -1:
seed_idx = max(available_segments, key=lambda x: -scores[x])
optimization_sign = -1
else:
raise ValueError(f"sigma must be -1, 1, or 'auto', got {sigma!r}")

return seed_idx, optimization_sign


def build_all_regions(
method: ExplanationMethodFn,
predictor: ModelPredictor,
input_batch: torch.Tensor,
replacement_image: torch.Tensor,
image_graph: ImageGraph,
target_class_idx: int,
original_log_odds: torch.Tensor,
scores: dict[int, float],
max_regions: int,
original_prob: float,
sigma: SeedSelectionMode = None,
desired_length: int = 30,
batch_size: int = 64,
) -> list[RegionResult]:
Expand All @@ -34,8 +65,13 @@ def build_all_regions(
replacement_image: Replacement tensor
image_graph: Graph representation of image segments
target_class_idx: Target class index
original_log_odds: Pre-computed unmasked target-class log-odds (scalar tensor)
scores: Base segment scores
max_regions: Maximum number of regions to construct
original_prob: Pre-computed unmasked probability for the target class
sigma: Seed selection mode. ``"None"`` picks max abs score and then inherits its sign. ``1`` chooses the
highest raw score (positive evidence), ``-1`` chooses the lowest
raw score (negative evidence).
desired_length: Target number of segments per region
batch_size: Batch size for model evaluation

Expand All @@ -45,8 +81,6 @@ def build_all_regions(
regions: list[RegionResult] = []
used_segments: set[int] = set()

original_prob = predictor.get_predictions(input_batch)[0, target_class_idx].item()

for _ in range(max_regions):
# Find best unprocessed seed
available_segments = [
Expand All @@ -56,9 +90,9 @@ def build_all_regions(
if not available_segments:
break

seed_idx = max(available_segments, key=lambda x: abs(scores[x]))
seed_score = scores[seed_idx]
optimization_sign = 1 if seed_score >= 0 else -1
seed_idx, optimization_sign = _select_seed_and_sign(
scores=scores, available_segments=available_segments, sigma=sigma
)

# Construct a SearchContext for the current step
ctx = SearchContext(
Expand All @@ -67,6 +101,7 @@ def build_all_regions(
replacement_image=replacement_image,
image_graph=image_graph,
target_class_idx=target_class_idx,
original_log_odds=original_log_odds,
seed_idx=seed_idx,
optimization_sign=optimization_sign,
used_segments=frozenset(used_segments),
Expand Down
1 change: 1 addition & 0 deletions ciao/algorithm/context.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ class SearchContext:
replacement_image: torch.Tensor
image_graph: ImageGraph
target_class_idx: int
original_log_odds: torch.Tensor
seed_idx: int
optimization_sign: int
used_segments: frozenset[int]
Expand Down
30 changes: 25 additions & 5 deletions ciao/algorithm/lookahead.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
Rolling horizon strategy: Look ahead multiple steps but only commit one step at a time.
"""

import time
from collections import deque
from collections.abc import Set

Expand Down Expand Up @@ -44,7 +45,9 @@ def build_region_greedy_lookahead(

eval_count = 0
trajectory: list[dict[str, float]] = []
current_best_score = -float("inf")
current_best_objective = -float("inf") # sign-normalized, for comparisons only
current_best_score = 0.0 # raw score, always set before first use
t0 = time.monotonic()

# Grow region one step at a time
while len(current_region) < desired_length:
Expand All @@ -71,6 +74,7 @@ def build_region_greedy_lookahead(
replacement_image=ctx.replacement_image,
target_class_idx=ctx.target_class_idx,
batch_size=ctx.batch_size,
original_log_odds=ctx.original_log_odds,
)
eval_count += len(candidate_regions)

Expand All @@ -83,9 +87,16 @@ def build_region_greedy_lookahead(
first_step = candidates[best_region]

signed_best = best_score * optimization_sign
if signed_best > current_best_score:
current_best_score = signed_best
trajectory.append({"evals": eval_count, "best_score": current_best_score})
if signed_best > current_best_objective:
current_best_objective = signed_best
current_best_score = best_score
trajectory.append(
{
"evals": eval_count,
"best_score": current_best_score,
"time": time.monotonic() - t0,
}
)
Comment thread
dhalmazna marked this conversation as resolved.

# Optimization - commit entire path
if len(best_region) == desired_length:
Expand All @@ -109,9 +120,18 @@ def build_region_greedy_lookahead(
replacement_image=ctx.replacement_image,
target_class_idx=ctx.target_class_idx,
batch_size=1,
original_log_odds=ctx.original_log_odds,
)[0]
eval_count += 1
trajectory.append({"evals": eval_count, "best_score": current_best_score})
if final_score * optimization_sign > current_best_objective:
current_best_score = final_score
trajectory.append(
{
"evals": eval_count,
"best_score": current_best_score,
"time": time.monotonic() - t0,
}
)

return RegionResult(
region=current_region,
Expand Down
Loading