Skip to content

API.md

Latest commit

 

History

History
62 lines (45 loc) · 1.29 KB

API.md

File metadata and controls

62 lines (45 loc) · 1.29 KB
Raw

API Notes

find_optimal_depth

find_optimal_depth(
    model,
    val_loader,
    task_type="classification",
    depths=None,
    device="cuda",
    num_epochs=10,
    lr=0.001,
    verbose=True,
    max_val_batches=None,
)

Runs a lightweight probe at each candidate depth and returns a DepthSearchResult.

Important notes:

  • val_loader is treated as a probe dataset in this package.
  • classification uses a linear classifier probe.
  • pose and detection use small regression-style probe heads.
  • custom currently falls back to the classification probe path.

FeatureExtractor

FeatureExtractor(model, layer_name, pooling="avg")

pooling options:

  • "avg": average-pool spatial or token dimensions
  • "max": max-pool spatial or token dimensions
  • "flatten": flatten everything after batch dimension
  • None: return raw activations

suggest_depths_for_task

Returns a small default sweep for:

  • classification
  • detection
  • pose
  • custom

get_available_layers

Discovers likely probe layers for common model layouts:

  • ResNet-style layer1 to layer4
  • stage-based stages.N
  • transformer-style blocks.N
  • feature-list features.N

CLI

  • orka-validate: installation smoke test
  • orka-find: synthetic probe run for fast validation