vertebrae.extractors.torch

Optional Torch module extractor for local user-supplied models.

Classes

TorchExtractor

Wrap a locally loaded PyTorch module as a feature extractor.

Module Contents

class vertebrae.extractors.torch.TorchExtractor(name, model, collate_fn, output_fn=None, outputs=None, device=None, modality='unknown', extractor_type='custom_torch', recipe_data=None, allow_sparse=False, streaming_safe=True, move_batch_to_device=True, move_model_to_device=True, spatial_output_fn=None, spatial_output_specs=None, structured_output_fn=None, structured_output_specs=None, checkpoint_paths=None, eval_mode=True, inference_mode=True, restore_model_mode=True, cache_identity=None)[source]

Wrap a locally loaded PyTorch module as a feature extractor.

Parameters:

  • name (str) – User-facing extractor name.

  • model (Any) – A locally loaded torch.nn.Module or compatible callable object.

  • collate_fn (Callable[[Any], Any]) – Callable that converts a batch of raw inputs into model inputs.

  • output_fn (Optional[Callable[[Any], Any]]) – Optional callable that converts raw model output into embeddings.

  • outputs (Optional[Sequence[Dict[str, Any]]]) – Optional named ordinary embedding-output declarations.

  • device (Optional[str]) – Optional device string passed to model.to(...) and batch items.

  • modality (str) – Input modality metadata.

  • extractor_type (str) – Extractor family metadata.

  • recipe_data (Optional[Dict[str, Any]]) – Extra serializable metadata for reproducibility.

  • allow_sparse (bool) – Whether sparse embedding outputs are accepted.

  • streaming_safe (bool) – Whether independent batches can be embedded without full-context state.

  • move_batch_to_device (bool) – Whether tensor-like batch items should be moved to device.

  • move_model_to_device (bool) – Whether the model should be moved to device once on first use.

  • spatial_output_fn (Optional[Callable[[Any], Any]])

  • spatial_output_specs (Optional[Iterable[vertebrae.extractors.spatial.SpatialOutputSpec]])

  • structured_output_fn (Optional[Callable[[Any], Any]])

  • structured_output_specs (Optional[Iterable[vertebrae.extractors.structured.StructuredOutputSpec]])

  • checkpoint_paths (Optional[Iterable[str]])

  • eval_mode (bool)

  • inference_mode (bool)

  • restore_model_mode (bool)

  • cache_identity (Optional[str])

get_resource_profile_adapter()[source]

Return Torch-specific synchronization and footprint hooks.

Return type:

Any

fit(X, y=None)[source]

No-op fit for local Torch models.

Parameters:

  • X (Any)

  • y (Any)

Return type:

TorchExtractor

transform(X)[source]

Apply the collate function, run the model, and validate embeddings.

Parameters:

X (Any)

Return type:

numpy.ndarray

output_specs()[source]

Return the declared ordinary embedding outputs.

Return type:

List[vertebrae.extractors.base.EmbeddingOutputSpec]

transform_many(X)[source]

Apply the model once and materialize every declared embedding output.

Parameters:

X (Any)

Return type:

List[vertebrae.extractors.base.EmbeddingOutput]

fit_transform(X, y=None)[source]

Fit the extractor and transform inputs.

Parameters:

  • X (Any)

  • y (Any)

Return type:

numpy.ndarray

recipe()[source]

Return a serializable recipe for this extractor.

Return type:

Dict[str, Any]