autogluon.multimodal 1.2.1b20250303__py3-none-any.whl → 1.2.1b20250305__py3-none-any.whl

autogluon/multimodal/models/__init__.py

@@ -1,4 +1,4 @@
-from . import utils
+from .augmenter import Augmenter
 from .categorical_mlp import CategoricalMLP
 from .clip import CLIPForImageText
 from .document_transformer import DocumentTransformer
@@ -9,7 +9,8 @@ from .fusion import (
     MultimodalFusionNER,
     MultimodalFusionTransformer,
 )
-from .huggingface_text import HFAutoModelForTextPrediction
+from .hf_text import HFAutoModelForTextPrediction
+from .meta_transformer import MetaTransformer
 from .mmdet_image import MMDetAutoModelForObjectDetection
 from .mmocr_text_detection import MMOCRAutoModelForTextDetection
 from .mmocr_text_recognition import MMOCRAutoModelForTextRecognition
@@ -18,4 +19,12 @@ from .numerical_mlp import NumericalMLP
 from .sam import SAMForSemanticSegmentation
 from .t_few import TFewModel
 from .timm_image import TimmAutoModelForImagePrediction
-from .utils import get_model_postprocess_fn
+from .utils import (
+    create_fusion_model,
+    create_model,
+    get_model_postprocess_fn,
+    is_lazy_weight_tensor,
+    list_timm_models,
+    modify_duplicate_model_names,
+    select_model,
+)

autogluon/multimodal/models/augmenter.py

@@ -0,0 +1,175 @@
+import logging
+
+import torch
+import torch.nn as nn
+from omegaconf import DictConfig
+from torch.nn import TransformerEncoder, TransformerEncoderLayer
+
+from .mlp import Unit
+
+logger = logging.getLogger(__name__)
+
+
+class VAETransformer(nn.Module):
+    def __init__(self, config: DictConfig, in_feautres: int, n_modality: int) -> None:
+        super().__init__()
+        self.config = config
+        self.emb_d = in_feautres
+        self.n_modality = n_modality
+        logger.debug(f" VAE Transformer # features {n_modality}, dim {self.emb_d}")
+
+        # encoder
+        encoder_layers = TransformerEncoderLayer(self.emb_d, config.n_head, config.tran_hidden, norm_first=True)
+        self.transformer_encoder = TransformerEncoder(encoder_layers, config.n_layer)
+
+        # encoder linear z
+        self.encoder_fc_z_mu = nn.Linear(self.emb_d, self.config.z_dim)
+        self.encoder_fc_z_logvar = nn.Linear(self.emb_d, self.config.z_dim)
+
+        # decoder linezr z
+        self.decoder_fc = nn.Linear(self.config.z_dim, self.emb_d)
+
+        # decoder
+        decoder_layers = TransformerEncoderLayer(self.emb_d, config.n_head, config.tran_hidden, norm_first=True)
+        self.transformer_decoder = TransformerEncoder(decoder_layers, config.n_layer)
+
+        self.last_layer = nn.Linear(self.emb_d, self.emb_d)
+
+        self.gating = nn.Identity()
+        self.init_parameters()
+
+    def init_parameters(self):
+        self.last_layer.weight.data.zero_()
+        self.last_layer.bias.data.zero_()
+
+    def reparameterize(self, mu, logvar):
+        std = torch.exp(0.5 * logvar)
+        eps = torch.randn_like(std)
+        return mu + eps * std
+
+    def forward(self, X):
+        input = X.reshape(-1, self.n_modality, self.emb_d)  # [B, # modality, emb dim] torch.Size([8, 3, 1024])
+
+        hidden = self.transformer_encoder(input)
+
+        z_mu, z_logvar = self.encoder_fc_z_mu(hidden), self.encoder_fc_z_logvar(hidden)
+
+        z = self.reparameterize(z_mu, z_logvar)
+
+        hidden = self.decoder_fc(z)
+
+        noise = self.gating(self.last_layer(self.transformer_decoder(hidden)[:, : self.n_modality, :]))
+        recon_x = X.reshape(-1, self.n_modality, self.emb_d) + noise
+
+        return recon_x.reshape(len(X), -1), z_mu, z_logvar
+
+
+class MlpVAE(nn.Module):
+    def __init__(self, input_dim, hidden_dim, z_dim=16) -> None:
+        super().__init__()
+        self.input_dim = input_dim
+        self.z_dim = z_dim
+        self.hidden_dim = hidden_dim
+
+        # Encoder P(Z|X)
+        encoder_layers = []
+        dims = [input_dim] + hidden_dim
+        for i in range(len(dims) - 1):
+            encoder_layers.append(
+                Unit(
+                    normalization="layer_norm",
+                    in_features=dims[i],
+                    out_features=dims[i + 1],
+                    activation="relu",
+                    dropout=0.5,
+                )
+            )
+        self.encoder = nn.Sequential(*encoder_layers)
+
+        self.encoder_fc_z_mu = nn.Linear(self.hidden_dim[-1], self.z_dim)
+        self.encoder_fc_z_logvar = nn.Linear(self.hidden_dim[-1], self.z_dim)
+
+        # Decoder P(X|Z)
+        decoder_layers = []
+        dims = [input_dim] + hidden_dim + [z_dim]
+
+        for i in range(len(dims) - 1, 0, -1):
+            decoder_layers.append(
+                Unit(
+                    normalization="layer_norm",
+                    in_features=dims[i],
+                    out_features=dims[i - 1],
+                    activation="relu",
+                    dropout=0.5,
+                )
+            )
+        self.decoder = nn.Sequential(*decoder_layers)
+
+        self.init_parameters()
+
+    def init_parameters(self):
+        self.decoder[-1].fc.weight.data.zero_()
+        self.decoder[-1].fc.bias.data.zero_()
+
+    def reparameterize(self, mu, logvar):
+        std = torch.exp(0.5 * logvar)
+        eps = torch.randn_like(std)
+        return mu + eps * std
+
+    def forward(self, x):
+        hidden = self.encoder(x)
+        z_mu, z_logvar = self.encoder_fc_z_mu(hidden), self.encoder_fc_z_logvar(hidden)
+        z = self.reparameterize(z_mu, z_logvar)
+
+        noise_x = self.decoder(z)
+        recon_x = x + noise_x
+        return recon_x, z_mu, z_logvar
+
+
+class Augmenter(nn.Module):
+    def __init__(
+        self,
+        arch_type: str,
+        input_dim: int,
+        z_dim: int,
+        num_layers: int,
+        adv_weight: float,
+    ) -> None:
+        super().__init__()
+        logger.debug("Initializing Augmenter")
+        self.arch_type = arch_type
+        self.input_dim = input_dim
+        self.z_dim = z_dim
+        self.num_layers = num_layers
+        self.adv_weight = adv_weight
+        logger.debug(f"augmenter arch_type: {self.arch_type}")
+        logger.debug(f"augmenter input_dim: {self.input_dim}")
+        logger.debug(f"augmenter z_dim: {self.z_dim}")
+        logger.debug(f"augmenter num_layers: {self.num_layers}")
+        logger.debug(f"augmenter adv_weight: {self.adv_weight}")
+        if self.arch_type == "mlp_vae":
+            step = int((self.input_dim - self.z_dim) / (self.num_layers + 1))
+            hidden = [*range(self.input_dim - step, self.z_dim + step, -step)]
+            self.vae = MlpVAE(input_dim=self.input_dim, hidden_dim=hidden, z_dim=self.z_dim)
+        else:
+            raise ValueError(f"Unknown arch_type: {self.arch_type}")
+
+        self.name_to_id = self.get_layer_ids()
+
+    def forward(self, x):
+        return self.vae(x)
+
+    def get_layer_ids(
+        self,
+    ):
+        """
+        All layers have the same id 0 since there is no pre-trained models used here.
+
+        Returns
+        -------
+        A dictionary mapping the layer names (keys) to their ids (values).
+        """
+        name_to_id = {}
+        for n, _ in self.named_parameters():
+            name_to_id[n] = 0
+        return name_to_id

autogluon/multimodal/models/categorical_mlp.py

@@ -1,4 +1,5 @@
-from typing import List, Optional
+import logging
+from typing import Dict, Optional
 
 import torch
 from torch import nn
@@ -7,6 +8,8 @@ from ..constants import CATEGORICAL, FEATURES, LABEL, LOGITS
 from .mlp import MLP
 from .utils import init_weights
 
+logger = logging.getLogger(__name__)
+
 
 class CategoricalMLP(nn.Module):
     """
@@ -17,11 +20,11 @@ class CategoricalMLP(nn.Module):
     def __init__(
         self,
         prefix: str,
-        num_categories: List[int],
+        num_categories: Dict,
         out_features: Optional[int] = None,
         num_layers: Optional[int] = 1,
         activation: Optional[str] = "gelu",
-        dropout_prob: Optional[float] = 0.5,
+        dropout: Optional[float] = 0.5,
         normalization: Optional[str] = "layer_norm",
         num_classes: Optional[int] = 0,
     ):
@@ -38,7 +41,7 @@ class CategoricalMLP(nn.Module):
             Number of MLP layers.
         activation
             Name of activation function.
-        dropout_prob
+        dropout
             Dropout probability.
         normalization
             Name of normalization function.
@@ -46,15 +49,17 @@ class CategoricalMLP(nn.Module):
             Number of classes. 1 for a regression task.
         """
         super().__init__()
+        logger.debug(f"initializing {prefix} (CategoricalMLP)")
         self.out_features = out_features
         max_embedding_dim = 100
         embed_exponent = 0.56
         size_factor = 1.0
         self.column_embeddings = nn.ModuleList()
         self.column_mlps = nn.ModuleList()
-        assert isinstance(num_categories, list)
+        assert isinstance(num_categories, dict)
+        self.num_categories = num_categories
 
-        for num_categories_per_col in num_categories:
+        for num_categories_per_col in num_categories.values():
             embedding_dim_per_col = int(
                 size_factor * max(2, min(max_embedding_dim, 1.6 * num_categories_per_col**embed_exponent))
             )
@@ -72,7 +77,7 @@ class CategoricalMLP(nn.Module):
                     out_features=out_features,
                     num_layers=num_layers,
                     activation=activation,
-                    dropout_prob=dropout_prob,
+                    dropout=dropout,
                     normalization=normalization,
                 )
             )
@@ -83,7 +88,7 @@ class CategoricalMLP(nn.Module):
             out_features=out_features,
             num_layers=num_layers,
             activation=activation,
-            dropout_prob=dropout_prob,
+            dropout=dropout,
             normalization=normalization,
         )
 

autogluon/multimodal/models/clip.py

@@ -5,7 +5,6 @@ import torch
 from torch import nn
 
 from ..constants import (
-    AUTOMM,
     COLUMN,
     COLUMN_FEATURES,
     FEATURES,
@@ -22,8 +21,12 @@ from .utils import (
     assign_layer_ids,
     get_column_features,
     get_hf_config_and_model,
+    get_image_size_mean_std,
     get_pretrained_tokenizer,
+    get_text_segment_num,
+    get_text_token_max_len,
     init_weights,
+    replace_missing_images_with_learnable,
 )
 
 logger = logging.getLogger(__name__)
@@ -42,6 +45,15 @@ class CLIPForImageText(nn.Module):
         num_classes: Optional[int] = None,
         pretrained: Optional[bool] = True,
         tokenizer_name: Optional[str] = "clip",
+        has_image: Optional[bool] = True,
+        has_text: Optional[bool] = True,
+        image_size: Optional[int] = None,
+        image_norm: Optional[str] = None,
+        image_chan_num: Optional[int] = 3,
+        use_learnable_image: Optional[bool] = False,
+        max_text_len: Optional[int] = None,
+        text_segment_num: Optional[int] = 1,
+        is_matching: Optional[bool] = False,
     ):
         """
         Load the pretrained CLIP from huggingface transformers.
@@ -60,16 +72,26 @@ class CLIPForImageText(nn.Module):
             Name of the huggingface tokenizer type.
         """
         super().__init__()
-        logger.debug(f"initializing {checkpoint_name}")
+        logger.debug(f"initializing {prefix} (CLIPForImageText)")
+        logger.debug(f"model checkpoint: {checkpoint_name}")
         self.checkpoint_name = checkpoint_name
         self.num_classes = num_classes
+        if is_matching:  # init both image and text attributes for matching
+            has_image, has_text = True, True
+        self.has_image = has_image
+        self.has_text = has_text
 
         self.config, self.model = get_hf_config_and_model(checkpoint_name=checkpoint_name, pretrained=pretrained)
-        self.tokenizer_name = tokenizer_name
-        self.tokenizer = get_pretrained_tokenizer(
-            tokenizer_name=self.tokenizer_name,
-            checkpoint_name=self.checkpoint_name,
-        )
+
+        if not self.has_image:
+            self.config.vision_config = None
+            self.model.vision_model = None
+            self.model.visual_projection = None
+
+        if not self.has_text:
+            self.config.text_config = None
+            self.model.text_model = None
+            self.model.text_projection = None
 
         self.out_features = self.model.config.projection_dim
 
@@ -77,6 +99,35 @@ class CLIPForImageText(nn.Module):
         self.head.apply(init_weights)
 
         self.prefix = prefix
+        if has_image:
+            self.image_size, self.image_mean, self.image_std = get_image_size_mean_std(
+                model_name=self.prefix,
+                config=self.model.vision_model.config,
+                provided_size=image_size,
+                provided_norm_type=image_norm,
+                support_variable_input_size=False,
+            )
+            self.use_learnable_image = use_learnable_image
+            if self.use_learnable_image:
+                self.learnable_image = nn.Parameter(torch.zeros(image_chan_num, self.image_size, self.image_size))
+                logger.debug("will use a learnable image to replace missing ones")
+        if has_text:
+            self.tokenizer_name = tokenizer_name
+            self.tokenizer = get_pretrained_tokenizer(
+                tokenizer_name=self.tokenizer_name,
+                checkpoint_name=self.checkpoint_name,
+            )
+            self.max_text_len = get_text_token_max_len(
+                provided_max_len=max_text_len,
+                config=self.model.text_model.config,
+                tokenizer=self.tokenizer,
+                checkpoint_name=self.checkpoint_name,
+            )
+            self.text_segment_num = get_text_segment_num(
+                config=self.model.text_model.config,
+                provided_segment_num=text_segment_num,
+                checkpoint_name=self.checkpoint_name,
+            )
 
         self.name_to_id = self.get_layer_ids()
         self.head_layer_names = [n for n, layer_id in self.name_to_id.items() if layer_id == 0]
@@ -117,6 +168,15 @@ class CLIPForImageText(nn.Module):
     def image_feature_dim(self):
         return self.model.config.vision_config.hidden_size
 
+    @property
+    def input_keys(self):
+        ret = []
+        if self.has_image:
+            ret.extend([self.image_key, self.image_valid_num_key])
+        if self.has_text:
+            ret.extend([self.text_token_ids_key, self.text_valid_length_key])
+        return ret
+
     def forward(
         self,
         batch: dict,
@@ -132,8 +192,8 @@ class CLIPForImageText(nn.Module):
         -------
             A dictionary with logits and features.
         """
-        has_image = self.image_key in batch
-        has_text = self.text_token_ids_key in batch
+        has_image = self.has_image and self.image_key in batch
+        has_text = self.has_text and self.text_token_ids_key in batch
         ret = {COLUMN_FEATURES: {FEATURES: {}, MASKS: {}}}
 
         if has_image:
@@ -141,6 +201,14 @@ class CLIPForImageText(nn.Module):
             image_valid_num = batch[self.image_valid_num_key]
             assert images.dim() == 5
             b, n, c, h, w = images.shape
+            steps = torch.arange(0, n).type_as(image_valid_num)
+            image_masks = steps.reshape((1, -1)) < image_valid_num.reshape((-1, 1))  # (b, n)
+            if self.use_learnable_image:
+                images = replace_missing_images_with_learnable(
+                    images=images,
+                    image_masks=image_masks,
+                    learnable_image=self.learnable_image,
+                )
             vision_outputs = self.model.vision_model(
                 pixel_values=images.reshape((b * n, c, h, w)),
                 output_attentions=True,
@@ -148,9 +216,9 @@ class CLIPForImageText(nn.Module):
                 return_dict=True,
             )
             image_features = self.model.visual_projection(vision_outputs.pooler_output)
-            steps = torch.arange(0, n).type_as(image_valid_num)
-            image_masks = (steps.reshape((1, -1)) < image_valid_num.reshape((-1, 1))).type_as(image_features)  # (b, n)
-            image_features = image_features.reshape((b, n, -1)) * image_masks[:, :, None]  # (b, n, num_features)
+            image_features = image_features.reshape((b, n, -1))  # (b, n, num_features)
+            if not self.use_learnable_image:
+                image_features = image_features * image_masks[:, :, None].type_as(image_features)
 
             # normalized features
             image_features = image_features / torch.clamp(image_features.norm(dim=-1, keepdim=True), min=1e-6)
@@ -199,18 +267,24 @@ class CLIPForImageText(nn.Module):
             ret[COLUMN_FEATURES][MASKS].update(text_column_feature_masks)
             ret[FEATURES] = text_features
 
-        if has_image and has_text:
-            if self.num_classes:
+        if self.num_classes:
+            if has_image and has_text:
                 features = image_features + text_features
                 logits = self.head(features)
                 ret[FEATURES] = features
+            elif has_image:
+                logits = self.head(image_features)
+            elif has_text:
+                logits = self.head(text_features)
             else:
+                raise RuntimeError("Neither image or text are used. Must have at least one.")
+            ret[LOGITS] = logits
+        else:
+            ret[LOGIT_SCALE] = self.model.logit_scale.exp()
+            if has_image and has_text:
                 # cosine similarity as logits
                 logits = torch.sum(image_features * text_features, dim=-1)
-
-            ret[LOGITS] = logits
-
-        ret[LOGIT_SCALE] = self.model.logit_scale.exp()
+                ret[LOGITS] = logits
 
         return {self.prefix: ret}