Module

embed.py

import goldener.embed
View source

Extraction of embeddings from models.

The embedding tools are used by the goldener.describe.GoldDescriptor to compute the embeddings of the samples:

On this page 6 classes
GoldEmbeddingTool EmbeddingFusionStrategy GoldTorchEmbeddingToolConfig GoldEmbeddingFusionTool GoldTorchEmbeddingTool GoldMultiModalTorchEmbeddingTool
class GoldEmbeddingTool(abc.ABC):

Abstract base class for embedding extraction from models.

This class defines the interface for embedding tools that can embed and optionally fuse embeddings from models. Implementations should provide specific mechanisms for extracting embeddings from different types of models (e.g., PyTorch, multimodal).

@abstractmethod
def embed(self, *args: Any, **kwargs: Any) -> dict[str, torch.Tensor]:

Embed the input data using the model.

Returns: Dictionary mapping layer names to their embedding tensors.

@abstractmethod
def embed_and_fuse(self, *args: Any, **kwargs: Any) -> torch.Tensor:

Embed and fuse the input data using the model.

Returns: Fused embedding tensor.

class EmbeddingFusionStrategy(enum.Enum):

Strategies to fuse embeddings from multiple layers.

CONCAT: Concatenate embeddings along the channel dimension. ADD: Element-wise addition of embeddings. AVERAGE: Element-wise average of embeddings. MAX: Element-wise maximum of embeddings.

CONCAT = <EmbeddingFusionStrategy.CONCAT: 'concat'>
AVERAGE = <EmbeddingFusionStrategy.AVERAGE: 'average'>
@dataclass
class GoldTorchEmbeddingToolConfig:

Configuration for the GoldTorchEmbeddingTool.

Attributes:
  • model: The PyTorch model from which to extract embeddings.
  • layers: List of layer names or a dictionary mapping group names to lists of layer names. If None, the last layer of the model is used.
  • channel_pos: Localization of the channel dimension in the embedding tensors (default is 1). It is used for concatenation and to determine the spatial dimensions for interpolation.
  • layer_fusion: Strategy to fuse embeddings from multiple layers within the same group.
  • group_fusion: Strategy to fuse embeddings from different groups.
GoldTorchEmbeddingToolConfig( model: torch.nn.modules.module.Module, layers: list[str] | dict[str, list[str]] | None = None, channel_pos: int = 1, layer_fusion: EmbeddingFusionStrategy = <EmbeddingFusionStrategy.CONCAT: 'concat'>, group_fusion: EmbeddingFusionStrategy = <EmbeddingFusionStrategy.CONCAT: 'concat'>)
model: torch.nn.modules.module.Module
layers: list[str] | dict[str, list[str]] | None = None
channel_pos: int = 1
class GoldEmbeddingFusionTool:

Embedding fusion from multiple layers and groups.

Attributes:
  • layer_fusion: Strategy to fuse embeddings from multiple layers within the same group.
  • group_fusion: Strategy to fuse embeddings from different groups.
GoldEmbeddingFusionTool( layer_fusion: EmbeddingFusionStrategy = <EmbeddingFusionStrategy.CONCAT: 'concat'>, group_fusion: EmbeddingFusionStrategy = <EmbeddingFusionStrategy.CONCAT: 'concat'>)

Initialize the GoldEmbeddingFusionTool.

Arguments:
  • layer_fusion: Strategy to fuse embeddings from multiple layers within the same group. Defaults to CONCAT.
  • group_fusion: Strategy to fuse embeddings from different groups. Defaults to CONCAT.
layer_fusion
group_fusion
@staticmethod
def fuse_tensors( tensors: List[torch.Tensor], strategy: EmbeddingFusionStrategy, channel_pos: int = 1) -> torch.Tensor:

Fuse a list of tensors.

The tensors are at least expected to have the shape (B, C), until (B, C, D, H, W) if sizes after the channel dimension differ. In this case, all tensors are interpolated to the largest size.

Arguments:
  • tensors: List of tensors to be fused.
  • strategy: Strategy to fuse the tensors.
  • channel_pos: Localization of the channel dimension in the embedding tensors (default is 1). It is used for concatenation and to determine the spatial dimensions for interpolation.

Returns: Fused tensors.

Raises:
  • ValueError: If the tensors have a different number of dimensions.
def fuse_embeddings( self, x: dict[str, torch.Tensor], layers: list[str] | dict[str, list[str]], channel_pos: int = 1) -> torch.Tensor:

Fuse embeddings from multiple layers and groups.

Arguments:
  • x: Dictionary mapping layer names to embedding tensors.
  • layers: List of layer names or a dictionary mapping group names to lists of layer names.
  • channel_pos: Localization of the channel dimension in the embedding tensors (default is 1). It is used for concatenation and to determine the spatial dimensions for interpolation.

Returns: Fused embedding tensor.

class GoldTorchEmbeddingTool(GoldEmbeddingTool):

Embedding tool for PyTorch models.

Once initialized, the tool registers forward hooks on the specified layers of the model. When the model processes input data, the hooks capture the outputs of these layers. The extracted embeddings can then be fused according to the specified strategies.

The model and layers cannot be changed after initialization. The embedding fusion can be changed.

Attributes:
  • _model: The PyTorch model from which to extract embeddings.
  • fusion: GoldEmbeddingFusionTool instance to handle embedding fusion.
  • _channel_pos: Localization of the channel dimension in the embedding tensors (default is 1). It is used for concatenation and to determine the spatial dimensions for interpolation.
  • _layers: List of layer names or a dictionary mapping group names to lists of layer names.
  • _hooks: Dictionary mapping layer names to their corresponding forward hook handles.
  • _embeddings: Dictionary to store extracted embeddings.
GoldTorchEmbeddingTool(config: GoldTorchEmbeddingToolConfig)

Initialize the GoldTorchEmbeddingTool.

Arguments:
  • config: Configuration object containing the model, layers, and fusion strategies.
fusion
model: torch.nn.modules.module.Module

The PyTorch model from which to extract embeddings.

layers: list[str] | dict[str, list[str]]

The layers from which to extract embeddings.

It also indicates the grouping of layers for fusion.

def embed(self, x: torch.Tensor) -> dict[str, torch.Tensor]:

Embed the input data using the model.

Arguments:
  • x: Input data tensor to be processed by the model.

Returns: Dictionary mapping layer names to their embedding tensors.

def embed_and_fuse(self, x: torch.Tensor) -> torch.Tensor:

Embed and fuse the input data using the model.

Arguments:
  • x: Input data tensor to be processed by the model.

Returns: Fused embedding tensor.

class GoldMultiModalTorchEmbeddingTool(GoldEmbeddingTool):

Embedding tool for multimodal data using PyTorch.

Each modality has its own GoldTorchEmbeddingTool defined by its own configuration. This allows for processing different types of input data (e.g., images, text, audio) with different models and then fusing their embeddings.

Attributes:
  • embedders: Dictionary mapping modality names to their GoldTorchEmbeddingTool instances.
  • strategy: Strategy for fusing embeddings from different modalities.
GoldMultiModalTorchEmbeddingTool( configs: Dict[str, GoldTorchEmbeddingToolConfig], strategy: EmbeddingFusionStrategy = <EmbeddingFusionStrategy.CONCAT: 'concat'>)

Initialize the multimodal embedding tool.

Arguments:
  • configs: Dictionary mapping modality names to their GoldTorchEmbeddingToolConfig.
  • strategy: Strategy to use for fusing embeddings from different modalities. Defaults to CONCAT.
embedders
strategy
def embed_and_fuse(self, x: Dict[str, torch.Tensor]) -> torch.Tensor:

Embed and fuse multimodal input data.

Arguments:
  • x: Dictionary mapping modality names to their input tensors.
Returns:

Fused embedding tensor combining all modalities.

def embed(self, x: Dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:

Embed multimodal input data without fusing.

Arguments:
  • x: Dictionary mapping modality names to their input tensors.
Returns:

Dictionary mapping "{modality}.{layer}" to their embedding tensors.