diff --git a/src/models/asymmetric_encoder.py b/src/models/asymmetric_encoder.py index 123f2d9..da78a1c 100644 --- a/src/models/asymmetric_encoder.py +++ b/src/models/asymmetric_encoder.py @@ -260,6 +260,9 @@ class TextFusionMLP(nn.Module): # Main model: AsymmetricEncoder # --------------------------------------------------------------------------- +# ResidualGateFusin experiment +from residual_fusions import ResidualGateType, GatedFusionResidual + class AsymmetricEncoder(nn.Module): """Dual encoder for CVGL with text fusion on both branches. @@ -294,6 +297,7 @@ class AsymmetricEncoder(nn.Module): def __init__( self, + gate_type: ResidualGateType = ResidualGateType.simple_residual_one_gate, dino_web_path: str = "nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth", dino_sat_path: str = "nn_models/DINO_SAT/model.safetensors", lrsclip_path: str = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt", @@ -365,8 +369,12 @@ class AsymmetricEncoder(nn.Module): ) # Separate gated fusion for query and gallery branches. - self.fusion_query = GatedFusion(init_gate=init_gate, baseline_mode=baseline_mode) - self.fusion_gallery = GatedFusion(init_gate=init_gate, baseline_mode=baseline_mode) + #! Experimental Gated fusion on query branch. + self.fusion_query = GatedFusionResidual(gate_type=gate_type, + init_gate=init_gate, baseline_mode=baseline_mode) + self.fusion_gallery = GatedFusionResidual(gate_type=gate_type, + init_gate=init_gate, baseline_mode=baseline_mode) + @staticmethod def _freeze(module: nn.Module) -> None: @@ -420,7 +428,7 @@ class AsymmetricEncoder(nn.Module): l1_texts: list[str] | None, l2_texts: list[str] | None, l3_texts: list[str] | None, - fusion: GatedFusion, + fusion: GatedFusionResidual, ) -> torch.Tensor: """Fuse image features with optional text, respecting per-sample presence. @@ -451,8 +459,7 @@ class AsymmetricEncoder(nn.Module): # Per-sample fusion: text-present samples use full gated fusion, # empty-caption samples pass through pure image features. - gate = torch.sigmoid(fusion.alpha) - fused_with_text = gate * img_feat + (1.0 - gate) * z_text + fused_with_text = fusion(img_feat, z_text) out = torch.where(has_text.unsqueeze(-1), fused_with_text, img_feat) return F.normalize(out, dim=-1) diff --git a/src/models/dual_encoder.py b/src/models/dual_encoder.py index 558d318..89d435e 100644 --- a/src/models/dual_encoder.py +++ b/src/models/dual_encoder.py @@ -21,6 +21,9 @@ import torch import torch.nn as nn import torch.nn.functional as F +# residual fusions exp +from residual_fusions import ResidualGateType, GatedFusionResidual + class ProjectionHead(nn.Module): """MLP projection head with L2 normalization.""" @@ -83,120 +86,6 @@ class GatedFusion(nn.Module): return torch.sigmoid(self.alpha).item() -#! GATE-FUSIONS MODIFICATIONS --------------------------------- -#! in_dim = 1024 - -from enum import Enum - - -class ResidualGateType(Enum): - simple_residual_one_gate = 0, - cross_gate = 1, - gate_sum = 2, - alpha_res_cat = 3, - alpha_res_sum = 4 - -# TODO: add GatedFusionresidual class to gin -@gin.configurable -class GatedFusionresidual(nn.Module): - """Learnable gated fusion of image and text embeddings. - - V1 - Simple residual gating with 1 common gate: - V2 - Cross residual gating with 2 cross-gates - V3 - Gate + Simple Sum of feats x & y - V4 - Alpha-weighted residual concat (per sample) - V5 - Alpha-weighted residual sum (per sample) - - """ - - def __init__(self, init_gate: float = 0.7, in_dim = 1024, baseline_mode: bool = False) -> None: - super().__init__() - # alpha is in logit space: sigmoid(alpha) = init_gate - init_alpha = torch.log(torch.tensor(init_gate / (1.0 - init_gate))) - self.alpha = nn.Parameter(init_alpha) - - # alphas for separated cases - init_alpha_img_cross_gate = torch.log(torch.tensor(init_gate / (1.0 - init_gate))) - init_alpha_text_cross_gate = torch.log(torch.tensor(init_gate / (1.0 - init_gate))) - self.alpha_img = nn.Parameter(init_alpha_img_cross_gate) - self.alpha_text = nn.Parameter(init_alpha_text_cross_gate) - - # weight for sum and cat residual - self.final_cat_residual_proj = nn.Linear(in_dim * 2, in_dim) - self.weight_net_for_sum = nn.Linear(in_dim, 1) - self.weight_net_for_cat = nn.Linear(in_dim * 2, 1) - - self.baseline_mode = baseline_mode - - def FuseSRGF(self, - img_feat: torch.Tensor, - text_feat: torch.Tensor | None, - ) -> torch.Tensor: - gate = torch.sigmoid(self.alpha) - img_res = img_feat * gate + img_feat - text_res = text_feat * (1 - gate) + text_feat - fused_vec = img_res + text_res - return fused_vec - - def FuseRCGF(self, - img_feat: torch.Tensor, - text_feat: torch.Tensor | None, - ) -> torch.Tensor: - gate_img = torch.sigmoid(self.alpha_img) - gate_text = torch.sigmoid(self.alpha_text) - z_img = img_feat + gate_text * img_feat - z_text = text_feat + gate_img * text_feat - fused_vec = z_img + z_text - return fused_vec - - def FuseGSUM(self, - img_feat: torch.Tensor, - text_feat: torch.Tensor | None, - ) -> torch.Tensor: - gate = torch.sigmoid(self.alpha) - fuzed_vec = img_feat + text_feat + gate * img_feat + (1.0 - gate) * text_feat - return fuzed_vec - - def FuseARGFSum( - self, - img_feat: torch.Tensor, - text_feat: torch.Tensor | None, - ) -> torch.Tensor: - gate = torch.sigmoid(self.alpha) - residual = img_feat + text_feat - res_weight = torch.sigmoid(self.weight_net_for_sum(residual)) - fuzed_vec = gate * img_feat + (1.0 - gate) * text_feat + res_weight * residual - return fuzed_vec - - def FuseARGFCat( - self, - img_feat: torch.Tensor, - text_feat: torch.Tensor | None, - ) -> torch.Tensor: - gate = torch.sigmoid(self.alpha) - cat_vec = torch.cat([img_feat, text_feat], dim=-1) - residual = self.final_cat_residual_proj(cat_vec) - res_weight = torch.sigmoid(self.weight_net_for_cat(cat_vec)) - fuzed_vec = gate * img_feat + (1.0 - gate) * text_feat + res_weight * residual - return fuzed_vec - - def forward( - self, - img_feat: torch.Tensor, - text_feat: torch.Tensor | None, - ) -> torch.Tensor: - if text_feat is None or self.baseline_mode: - return img_feat - gate = torch.sigmoid(self.alpha) - - # TODO: switch forwards here - - return gate * img_feat + (1.0 - gate) * text_feat - - - - - @gin.configurable class DualEncoderCaptionTest(nn.Module): """GeoRSCLIP dual encoder with gated text fusion on query branch. @@ -222,6 +111,7 @@ class DualEncoderCaptionTest(nn.Module): baseline_mode: bool = False, init_gate: float = 0.7, device: str = "cuda", + gate_type: ResidualGateType = ResidualGateType.simple_residual_one_gate ) -> None: super().__init__() self.embed_dim = embed_dim @@ -247,7 +137,11 @@ class DualEncoderCaptionTest(nn.Module): self._apply_unfreeze(unfreeze_mode) # Gated fusion on query branch. - self.fusion = GatedFusion(init_gate=init_gate, baseline_mode=baseline_mode) + # self.fusion = GatedFusion(init_gate=init_gate, baseline_mode=baseline_mode) + + #! Experimental Gated fusion on query branch. + self.fusion = GatedFusionResidual(gate_type=gate_type, + init_gate=init_gate, baseline_mode=baseline_mode) # Projection heads. self.proj_query = ProjectionHead( diff --git a/src/models/residual_fusions.py b/src/models/residual_fusions.py new file mode 100644 index 0000000..ff5c592 --- /dev/null +++ b/src/models/residual_fusions.py @@ -0,0 +1,144 @@ +import torch +import torch.nn as nn +import gin + +#! GATE-FUSIONS MODIFICATIONS --------------------------------- +#! in_dim = 1024 + +from enum import Enum +import math + +class ResidualGateType(Enum): + simple_residual_one_gate = 0, + cross_gate = 1, + gate_sum = 2, + alpha_res_cat = 3, + alpha_res_sum = 4 + +# TODO: add GatedFusionresidual class to gin + +def init_bias_for_sigmoid(linear: nn.Linear, value: float) -> None: + nn.init.zeros_(linear.weight) + nn.init.constant_(linear.bias, math.log(value / (1.0 - value))) + +def init_residual_projs(linear: nn.Linear, scale: float) -> None: + nn.init.xavier_uniform_(linear.weight, gain=scale) + nn.init.zeros_(linear.bias) + +RESIDUAL_GATES = { + ResidualGateType.alpha_res_sum, + ResidualGateType.alpha_res_cat +} + +@gin.configurable +class GatedFusionResidual(nn.Module): + """Learnable gated fusion of image and text embeddings. + + V1 - Simple residual gating with 1 common gate: + V2 - Cross residual gating with 2 cross-gates + V3 - Gate + Simple Sum of feats x & y + V4 - Alpha-weighted residual sum (per sample) + V5 - Alpha-weighted residual concat (per sample) + """ + + def __init__(self, gate_type: ResidualGateType, + init_gate: float = 0.7, in_dim = 1024, + init_res_weight: float = 0.1, residual_proj_scale: float = 0.1, + baseline_mode: bool = False, + ) -> None: + super().__init__() + # alpha is in logit space: sigmoid(alpha) = init_gate + init_alpha = torch.log(torch.tensor(init_gate / (1.0 - init_gate))) + self.alpha = nn.Parameter(init_alpha) + + # alphas for separated cases + if gate_type == ResidualGateType.cross_gate: + init_alpha_img_cross_gate = torch.log(torch.tensor(init_gate / (1.0 - init_gate))) + init_alpha_text_cross_gate = torch.log(torch.tensor(init_gate / (1.0 - init_gate))) + self.alpha_img = nn.Parameter(init_alpha_img_cross_gate) + self.alpha_text = nn.Parameter(init_alpha_text_cross_gate) + + # weight for sum and cat residual + if gate_type in RESIDUAL_GATES: + self.final_cat_residual_proj = nn.Linear(in_dim * 2, in_dim) + self.weight_net_for_sum = nn.Linear(in_dim, 1) + self.weight_net_for_cat = nn.Linear(in_dim * 2, 1) + init_bias_for_sigmoid(self.weight_net_for_sum, value=init_res_weight) + init_bias_for_sigmoid(self.weight_net_for_cat, value=init_res_weight) + init_residual_projs(self.final_cat_residual_proj, scale=residual_proj_scale) + + self.gate_type = gate_type + self.baseline_mode = baseline_mode + + def FuseSRGF(self, + img_feat: torch.Tensor, + text_feat: torch.Tensor | None, + ) -> torch.Tensor: + gate = torch.sigmoid(self.alpha) + img_res = img_feat * gate + img_feat + text_res = text_feat * (1 - gate) + text_feat + fused_vec = img_res + text_res + return fused_vec + + def FuseRCGF(self, + img_feat: torch.Tensor, + text_feat: torch.Tensor | None, + ) -> torch.Tensor: + gate_img = torch.sigmoid(self.alpha_img) + gate_text = torch.sigmoid(self.alpha_text) + z_img = img_feat + gate_text * img_feat + z_text = text_feat + gate_img * text_feat + fused_vec = z_img + z_text + return fused_vec + + def FuseGSUM(self, + img_feat: torch.Tensor, + text_feat: torch.Tensor | None, + ) -> torch.Tensor: + gate = torch.sigmoid(self.alpha) + fuzed_vec = img_feat + text_feat + gate * img_feat + (1.0 - gate) * text_feat + return fuzed_vec + + def FuseARGFSum( + self, + img_feat: torch.Tensor, + text_feat: torch.Tensor | None, + ) -> torch.Tensor: + gate = torch.sigmoid(self.alpha) + residual = img_feat + text_feat + res_weight = torch.sigmoid(self.weight_net_for_sum(residual)) + fuzed_vec = gate * img_feat + (1.0 - gate) * text_feat + res_weight * residual + return fuzed_vec + + def FuseARGFCat( + self, + img_feat: torch.Tensor, + text_feat: torch.Tensor | None, + ) -> torch.Tensor: + gate = torch.sigmoid(self.alpha) + cat_vec = torch.cat([img_feat, text_feat], dim=-1) + residual = self.final_cat_residual_proj(cat_vec) + res_weight = torch.sigmoid(self.weight_net_for_cat(cat_vec)) + fuzed_vec = gate * img_feat + (1.0 - gate) * text_feat + res_weight * residual + return fuzed_vec + + def forward( + self, + img_feat: torch.Tensor, + text_feat: torch.Tensor | None, + ) -> torch.Tensor: + if text_feat is None or self.baseline_mode: + return img_feat + + if self.gate_type == ResidualGateType.simple_residual_one_gate: + fused_vec = self.FuseSRGF(img_feat=img_feat, text_feat=text_feat) + if self.gate_type == ResidualGateType.cross_gate: + fused_vec = self.FuseRCGF(img_feat=img_feat, text_feat=text_feat) + if self.gate_type == ResidualGateType.gate_sum: + fused_vec = self.FuseGSUM(img_feat=img_feat, text_feat=text_feat) + if self.gate_type == ResidualGateType.alpha_res_sum: + fused_vec = self.FuseARGFSum(img_feat=img_feat, text_feat=text_feat) + if self.gate_type == ResidualGateType.alpha_res_cat: + fused_vec = self.FuseARGFCat(img_feat=img_feat, text_feat=text_feat) + + return fused_vec diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py index 24ae3ea..7e12c3e 100644 --- a/src/training/train_gtauav.py +++ b/src/training/train_gtauav.py @@ -99,7 +99,8 @@ class TrainConfigGTAUAV: # Training. resume_from: str | None = None # path to checkpoint for resuming - output_dir: str = "out/gtauav/with_text" + # output_dir: str = "out/gtauav/with_text" + output_dir: str = "out/gtauav/with_text_exp_gate_SRGF" epochs: int = 10 batch_size: int = 8 num_workers: int = 4