Uncertain fuse

Attn fuse
2026-02-01 17:48:47 +08:00 · 2026-01-31 23:48:05 +08:00
10 changed files with 133 additions and 53 deletions
--- a/MSGCoOp/README.md
+++ b/MSGCoOp/README.md
@@ -35,12 +35,12 @@ Follow [DATASETS.md](DATASETS.md) to install the datasets.

 ## Generalization From Base to New Classes

-You will need `base2new_train.sh`, `base2new_test.sh`, and `base2new_all.sh`. The scripts with the prefix `base2new_train` train a model on base classes while the ones with the prefix `base2new_test` evaluate the trained model on new classes. Both kinds of scripts have three input argument, i.e., `TRAINER SG_WEIGHT DIV_WEIGHT`.
+You will need `base2new_train.sh`, `base2new_test.sh`, and `base2new_all.sh`. The scripts with the prefix `base2new_train` train a model on base classes while the ones with the prefix `base2new_test` evaluate the trained model on new classes. Both kinds of scripts have three input argument, i.e., `TRAINER SG_WEIGHT DIV_WEIGHT ATTN_REG_WEIGHT UNCERTAINTY_SCALE`.

 You can run base to new on all datasets as follow:

 ```bash
-bash scripts/base2new_all.sh MSGCoOp 8.0 1.0
+bash scripts/base2new_all.sh MSGCoOp 8.0 1.0 0.01 0.5
 ```

 When the evaluation is done, you can use `extract_acc.py` (replace the `root_dir` in the `main` function to your output dir) to automatically calculate the average results. For instance, after you finish the trainning using the aforementioned commands, you would get
@@ -91,13 +91,13 @@ Then, you will get the avarage accuracy.
 Fisrt, you need train on all classes over ImageNet:

 ```bash
-bash scripts/xd_train.sh MSGCoOp 8.0 1.0
+bash scripts/xd_train.sh MSGCoOp 8.0 1.0 0.01 0.5
 ```

 Then you can evaluate the performance on other ImageNet variants by run:

 ```bash
-bash scripts/xdo_test.sh MSGCoOp 8.0 1.0
+bash scripts/xdo_test.sh MSGCoOp 8.0 1.0 0.01 0.5
 ```

 And you will get the `output_xdo` after script finish. You can get the accuracy by `extract_acc.py` (need modify the `root_dir` to `output_xdo` ).
--- a/MSGCoOp/configs/trainers/MSGCoOp/vit_b16_ep100_ctxv1.yaml
+++ b/MSGCoOp/configs/trainers/MSGCoOp/vit_b16_ep100_ctxv1.yaml
@@ -31,3 +31,5 @@ MODEL:
 TRAINER:
  COOP:
    CTX_INIT: True
+    ATTENTION_REG_WEIGHT: 0.01
+    UNCERTAINTY_SCALE: 0.5
--- a/MSGCoOp/extract_acc.py
+++ b/MSGCoOp/extract_acc.py
@@ -98,7 +98,7 @@ def print_model_results(results, model_name):
        print("No complete dataset results found for this model.")

 def main():
-    root_dir = 'output_xda'  # 修改为你的output目录路径
+    root_dir = 'output'  # 修改为你的output目录路径
    target_model = 'MSGCoOp'  # 指定要分析的模型
    
    results = collect_model_results(root_dir, target_model)
--- a/MSGCoOp/scripts/base2new_all.sh
+++ b/MSGCoOp/scripts/base2new_all.sh
@@ -3,36 +3,27 @@
 TRAINER=$1
 KG_WEIGHT=$2
 MP_WEIGHT=$3
+ATTN_REG_WEIGHT=$4
+UNCERTAINTY_SCALE=$5

-CUDA_VISIBLE_DEVICES=0 bash scripts/base2new_train.sh ${TRAINER} ucf101 ${KG_WEIGHT} ${MP_WEIGHT}
-CUDA_VISIBLE_DEVICES=0 bash scripts/base2new_test.sh ${TRAINER} ucf101 ${KG_WEIGHT} ${MP_WEIGHT}
+# Define datasets array
+datasets=(
+    "ucf101"
+    "eurosat"
+    "oxford_pets"
+    "food101"
+    "oxford_flowers"
+    "dtd"
+    "caltech101"
+    "fgvc_aircraft"
+    "stanford_cars"
+    "sun397"
+    "imagenet"
+)

-CUDA_VISIBLE_DEVICES=0 bash scripts/base2new_train.sh ${TRAINER} eurosat ${KG_WEIGHT} ${MP_WEIGHT}
-CUDA_VISIBLE_DEVICES=0 bash scripts/base2new_test.sh ${TRAINER} eurosat ${KG_WEIGHT} ${MP_WEIGHT}
+# Loop through datasets
+for dataset in "${datasets[@]}"; do
+    CUDA_VISIBLE_DEVICES=0 bash scripts/base2new_train.sh ${TRAINER} ${dataset} ${KG_WEIGHT} ${MP_WEIGHT} ${ATTN_REG_WEIGHT} ${UNCERTAINTY_SCALE}
+    CUDA_VISIBLE_DEVICES=0 bash scripts/base2new_test.sh ${TRAINER} ${dataset} ${KG_WEIGHT} ${MP_WEIGHT} ${ATTN_REG_WEIGHT} ${UNCERTAINTY_SCALE}
+done

-CUDA_VISIBLE_DEVICES=0 bash scripts/base2new_train.sh ${TRAINER} oxford_pets ${KG_WEIGHT} ${MP_WEIGHT}
-CUDA_VISIBLE_DEVICES=0 bash scripts/base2new_test.sh ${TRAINER} oxford_pets ${KG_WEIGHT} ${MP_WEIGHT}
-
-CUDA_VISIBLE_DEVICES=0 bash scripts/base2new_train.sh ${TRAINER} food101 ${KG_WEIGHT} ${MP_WEIGHT}
-CUDA_VISIBLE_DEVICES=0 bash scripts/base2new_test.sh ${TRAINER} food101 ${KG_WEIGHT} ${MP_WEIGHT}
-
-CUDA_VISIBLE_DEVICES=0 bash scripts/base2new_train.sh ${TRAINER} oxford_flowers ${KG_WEIGHT} ${MP_WEIGHT}
-CUDA_VISIBLE_DEVICES=0 bash scripts/base2new_test.sh ${TRAINER} oxford_flowers ${KG_WEIGHT} ${MP_WEIGHT}
-
-CUDA_VISIBLE_DEVICES=0 bash scripts/base2new_train.sh ${TRAINER} dtd ${KG_WEIGHT} ${MP_WEIGHT}
-CUDA_VISIBLE_DEVICES=0 bash scripts/base2new_test.sh ${TRAINER} dtd ${KG_WEIGHT} ${MP_WEIGHT}
-
-CUDA_VISIBLE_DEVICES=0 bash scripts/base2new_train.sh ${TRAINER} caltech101 ${KG_WEIGHT} ${MP_WEIGHT}
-CUDA_VISIBLE_DEVICES=0 bash scripts/base2new_test.sh ${TRAINER} caltech101 ${KG_WEIGHT} ${MP_WEIGHT}
-
-CUDA_VISIBLE_DEVICES=0 bash scripts/base2new_train.sh ${TRAINER} fgvc_aircraft ${KG_WEIGHT} ${MP_WEIGHT}
-CUDA_VISIBLE_DEVICES=0 bash scripts/base2new_test.sh ${TRAINER} fgvc_aircraft ${KG_WEIGHT} ${MP_WEIGHT}
-
-CUDA_VISIBLE_DEVICES=0 bash scripts/base2new_train.sh ${TRAINER} stanford_cars ${KG_WEIGHT} ${MP_WEIGHT}
-CUDA_VISIBLE_DEVICES=0 bash scripts/base2new_test.sh ${TRAINER} stanford_cars ${KG_WEIGHT} ${MP_WEIGHT}
-
-CUDA_VISIBLE_DEVICES=0 bash scripts/base2new_train.sh ${TRAINER} sun397 ${KG_WEIGHT} ${MP_WEIGHT}
-CUDA_VISIBLE_DEVICES=0 bash scripts/base2new_test.sh ${TRAINER} sun397 ${KG_WEIGHT} ${MP_WEIGHT}
-
-CUDA_VISIBLE_DEVICES=0 bash scripts/base2new_train.sh ${TRAINER} imagenet ${KG_WEIGHT} ${MP_WEIGHT}
-CUDA_VISIBLE_DEVICES=0 bash scripts/base2new_test.sh ${TRAINER} imagenet ${KG_WEIGHT} ${MP_WEIGHT}
--- a/MSGCoOp/scripts/base2new_test.sh
+++ b/MSGCoOp/scripts/base2new_test.sh
@@ -7,6 +7,8 @@ DATASET=$2
 N_PROMPTS=4
 KG_WEIGHT=$3
 MP_WEIGHT=$4
+ATTN_REG_WEIGHT=$5
+UNCERTAINTY_SCALE=$6
 #CFG=rn50_ep100  # config file
 CFG=vit_b16_ep100_ctxv1
 CTP=end  # class token position (end or middle)
@@ -19,7 +21,7 @@ SUB=new

 for SEED in 1 2 3
 do
-    COMMON_DIR=${DATASET}/shots_${SHOTS}_${KG_WEIGHT}/${TRAINER}/${CFG}/seed${SEED}
+    COMMON_DIR=${DATASET}/shots_${SHOTS}_${KG_WEIGHT}_${MP_WEIGHT}_${ATTN_REG_WEIGHT}_${UNCERTAINTY_SCALE}/${TRAINER}/${CFG}/seed${SEED}
    MODEL_DIR=output/base2new/train_base/${COMMON_DIR}
    DIR=output/base2new/test_${SUB}/${COMMON_DIR}

@@ -43,6 +45,7 @@ do
        TRAINER.COOP.CSC ${CSC} \
        TRAINER.COOP.CLASS_TOKEN_POSITION ${CTP} \
        DATASET.NUM_SHOTS ${SHOTS} \
-        DATASET.SUBSAMPLE_CLASSES ${SUB} 
+        DATASET.SUBSAMPLE_CLASSES ${SUB} \
+        TRAINER.COOP.UNCERTAINTY_SCALE ${UNCERTAINTY_SCALE} 
    fi
 done
--- a/MSGCoOp/scripts/base2new_train.sh
+++ b/MSGCoOp/scripts/base2new_train.sh
@@ -6,6 +6,8 @@ TRAINER=$1
 DATASET=$2
 KG_WEIGHT=$3
 MP_WEIGHT=$4
+ATTN_REG_WEIGHT=$5
+UNCERTAINTY_SCALE=$6
 N_PROMPTS=4
 #CFG=rn50_ep100  # config file\
 CFG=vit_b16_ep100_ctxv1
@@ -16,7 +18,7 @@ CSC=False  # class-specific context (False or True)

 for SEED in 1 2 3
 do
-    DIR=output/base2new/train_base/${DATASET}/shots_${SHOTS}_${KG_WEIGHT}/${TRAINER}/${CFG}/seed${SEED}
+    DIR=output/base2new/train_base/${DATASET}/shots_${SHOTS}_${KG_WEIGHT}_${MP_WEIGHT}_${ATTN_REG_WEIGHT}_${UNCERTAINTY_SCALE}/${TRAINER}/${CFG}/seed${SEED}
    if [ -d "$DIR" ]; then
        echo "Results are available in ${DIR}. Skip this job"
    else
@@ -35,6 +37,8 @@ do
        DATASET.NUM_SHOTS ${SHOTS} \
 	    DATASET.SUBSAMPLE_CLASSES base \
 	    TRAINER.COOP.N_PROMPTS ${N_PROMPTS} \
-        TRAINER.COOP.DIV_WEIGHT ${MP_WEIGHT}
+        TRAINER.COOP.DIV_WEIGHT ${MP_WEIGHT} \
+        TRAINER.COOP.ATTENTION_REG_WEIGHT ${ATTN_REG_WEIGHT} \
+        TRAINER.COOP.UNCERTAINTY_SCALE ${UNCERTAINTY_SCALE}
    fi
 done
--- a/MSGCoOp/scripts/xd_train.sh
+++ b/MSGCoOp/scripts/xd_train.sh
@@ -6,6 +6,8 @@ TRAINER=$1
 N_PROMPTS=3
 KG_WEIGHT=$2
 MP_WEIGHT=$3
+ATTN_REG_WEIGHT=$4
+UNCERTAINTY_SCALE=$5
 CFG=vit_b16_ep100_ctxv1 
 CTP=end  # class token position (end or middle)
 NCTX=4  # number of context tokens
@@ -18,7 +20,7 @@ for DATASET in ${SRC_DATASETS}
 do
    for SEED in 1 2 3
    do
-        DIR=output_xd/base2new/train_base/${DATASET}/shots_${SHOTS}_${KG_WEIGHT}/${TRAINER}/${CFG}/seed${SEED}
+        DIR=output_xd/base2new/train_base/${DATASET}/shots_${SHOTS}_${KG_WEIGHT}_${MP_WEIGHT}_${ATTN_REG_WEIGHT}_${UNCERTAINTY_SCALE}/${TRAINER}/${CFG}/seed${SEED}
        if [ -d "$DIR" ]; then
            echo "Results are available in ${DIR}. Skip this job"
        else
@@ -36,7 +38,9 @@ do
            TRAINER.COOP.CLASS_TOKEN_POSITION ${CTP} \
            DATASET.NUM_SHOTS ${SHOTS} \
            TRAINER.COOP.N_PROMPTS ${N_PROMPTS} \
-            TRAINER.COOP.DIV_WEIGHT ${MP_WEIGHT}
+            TRAINER.COOP.DIV_WEIGHT ${MP_WEIGHT} \
+            TRAINER.COOP.ATTENTION_REG_WEIGHT ${ATTN_REG_WEIGHT} \
+            TRAINER.COOP.UNCERTAINTY_SCALE ${UNCERTAINTY_SCALE}
        fi
    done
 done
--- a/MSGCoOp/scripts/xdo_test.sh
+++ b/MSGCoOp/scripts/xdo_test.sh
@@ -6,6 +6,8 @@ TRAINER=$1
 N_PROMPTS=3
 KG_WEIGHT=$2
 MP_WEIGHT=$3
+ATTN_REG_WEIGHT=$4
+UNCERTAINTY_SCALE=$5
 CFG=vit_b16_ep100_ctxv1 
 CTP=end  # class token position (end or middle)
 NCTX=4  # number of context tokens
@@ -19,8 +21,8 @@ for DATASET in imagenetv2 imagenet_sketch imagenet_a imagenet_r
 do
    for SEED in 1 2 3
    do
-        MODEL_DIR=output_xd/base2new/train_base/${SRC_DATASETS}/shots_${SHOTS}_${KG_WEIGHT}/${TRAINER}/${CFG}/seed${SEED}
-        DIR=output_xdo/base2new/train_base/${DATASET}/shots_${SHOTS}_${KG_WEIGHT}/${TRAINER}/${CFG}/seed${SEED}
+        MODEL_DIR=output_xd/base2new/train_base/${SRC_DATASETS}/shots_${SHOTS}_${KG_WEIGHT}_${MP_WEIGHT}_${ATTN_REG_WEIGHT}_${UNCERTAINTY_SCALE}/${TRAINER}/${CFG}/seed${SEED}
+        DIR=output_xdo/base2new/train_base/${DATASET}/shots_${SHOTS}_${KG_WEIGHT}_${MP_WEIGHT}_${ATTN_REG_WEIGHT}_${UNCERTAINTY_SCALE}/${TRAINER}/${CFG}/seed${SEED}
        if [ -d "$DIR" ]; then
            echo "Results are available in ${DIR}. Skip this job"
        else
@@ -41,7 +43,9 @@ do
            TRAINER.COOP.CLASS_TOKEN_POSITION ${CTP} \
            DATASET.NUM_SHOTS ${SHOTS} \
            TRAINER.COOP.N_PROMPTS ${N_PROMPTS} \
-            TRAINER.COOP.DIV_WEIGHT ${MP_WEIGHT}
+            TRAINER.COOP.DIV_WEIGHT ${MP_WEIGHT} \
+            TRAINER.COOP.ATTENTION_REG_WEIGHT ${ATTN_REG_WEIGHT} \
+            TRAINER.COOP.UNCERTAINTY_SCALE ${UNCERTAINTY_SCALE}
        fi
    done
 done
--- a/MSGCoOp/train.py
+++ b/MSGCoOp/train.py
@@ -105,6 +105,8 @@ def extend_cfg(cfg):
    cfg.TRAINER.COCOOP.PREC = "fp16"  # fp16, fp32, amp
    cfg.TRAINER.COOP.DIV_WEIGHT = 0.1
    cfg.TRAINER.COOP.N_PROMPTS = 3
+    cfg.TRAINER.COOP.ATTENTION_REG_WEIGHT = 0.01
+    cfg.TRAINER.COOP.UNCERTAINTY_SCALE = 0.5

    cfg.DATASET.SUBSAMPLE_CLASSES = "all"  # all, base or new
    """
--- a/MSGCoOp/trainers/msgcoop.py
+++ b/MSGCoOp/trainers/msgcoop.py
@@ -223,6 +223,58 @@ class Adapter(nn.Module):
        x = self.fc(x)
        return x

+class AttentionBasedIntegrator(nn.Module):
+    def __init__(self, img_dim=512, n_prompts=4, dtype=None, uncertainty_scale=0.5):
+        super().__init__()
+        self.attention = nn.Sequential(
+            nn.Linear(img_dim, img_dim // 4),
+            nn.Tanh(),
+            nn.Linear(img_dim // 4, n_prompts)
+        )
+        self.dtype = dtype
+        self.uncertainty_scale = uncertainty_scale
+        
+        if dtype is not None:
+            self.attention = self.attention.to(dtype)
+        
+    def forward(self, image_features, all_logits):
+        batch_size = image_features.shape[0]
+        n_prompts = all_logits.shape[0]
+        
+        # 注意力权重
+        attn_scores = self.attention(image_features)
+        attn_weights = F.softmax(attn_scores, dim=-1)  # [batch, n_prompts]
+        
+        # 不确定性权重（基于预测熵）- 使用log-sum-exp技巧提高稳定性
+        probs = F.softmax(all_logits, dim=-1)  # [n_prompts, batch, n_classes]
+        log_probs = F.log_softmax(all_logits, dim=-1)
+        
+        # 计算熵：H = -sum(p * log(p))
+        entropy = -(probs * log_probs).sum(dim=-1)  # [n_prompts, batch]
+        
+        # 转换为确定性（熵越小，确定性越大）
+        certainty = -entropy.t()  # [batch, n_prompts]
+        
+        # 归一化不确定性权重（使用log-sum-exp提高稳定性）
+        certainty_scaled = certainty * self.uncertainty_scale
+        log_uncertainty_weights = certainty_scaled - torch.logsumexp(certainty_scaled, dim=-1, keepdim=True)
+        uncertainty_weights = torch.exp(log_uncertainty_weights)
+        
+        # 混合权重：结合注意力权重和不确定性权重（使用对数域）
+        log_attn_weights = torch.log(attn_weights + 1e-8)
+        log_uncertainty_weights = torch.log(uncertainty_weights + 1e-8)
+        
+        # 在对数域中相加，然后指数化
+        log_hybrid_weights = torch.log(torch.tensor(0.5, device=attn_weights.device)) + \
+                             log_attn_weights + log_uncertainty_weights
+        log_hybrid_weights = log_hybrid_weights - torch.logsumexp(log_hybrid_weights, dim=-1, keepdim=True)
+        hybrid_weights = torch.exp(log_hybrid_weights)
+        
+        # 加权集成
+        weighted_logits = torch.einsum('bp,pbc->bc', hybrid_weights, all_logits)
+        
+        return weighted_logits, hybrid_weights, entropy
+
 class CustomCLIP(nn.Module):
    def __init__(self, cfg, classnames, clip_model):
        super().__init__()
@@ -237,6 +289,15 @@ class CustomCLIP(nn.Module):
        self.meta_net = self.prompt_learner.meta_net
        self.adapter = Adapter(512, 4).to(clip_model.dtype)
        
+        uncertainty_scale = getattr(cfg.TRAINER.COOP, 'UNCERTAINTY_SCALE', 0.5)
+        
+        self.prompt_integrator = AttentionBasedIntegrator(
+            img_dim=clip_model.visual.output_dim,
+            n_prompts=self.n_prompts,
+            dtype=clip_model.dtype,
+            uncertainty_scale=uncertainty_scale
+        )
+
    def compute_diversity_loss(self, text_features):
        if self.n_prompts == 1:
            return torch.tensor(0.0, device=text_features.device)
@@ -284,9 +345,11 @@ class CustomCLIP(nn.Module):
            logits_i = logit_scale * image_features @ text_features_i.t()
            all_logits.append(logits_i)
        
-        logits = torch.stack(all_logits).mean(dim=0)
+        all_logits = torch.stack(all_logits)
        
-        return logits, score, diversity_loss
+        logits, hybrid_weights, entropy = self.prompt_integrator(image_features, all_logits)
+
+        return logits, score, diversity_loss, hybrid_weights, entropy


@TRAINER_REGISTRY.register()
@@ -310,10 +373,11 @@ class MSGCoOp(TrainerX):
        self.model = CustomCLIP(cfg, classnames, clip_model)
        self.w = cfg.TRAINER.COOP.W
        self.diversity_weight = cfg.TRAINER.COOP.DIV_WEIGHT
+        self.attn_reg_weight = cfg.TRAINER.COOP.ATTENTION_REG_WEIGHT if hasattr(cfg.TRAINER.COOP, 'ATTENTION_REG_WEIGHT') else 0.01

        print("Turning off gradients in both the image and the text encoder")
        for name, param in self.model.named_parameters():
-            if "ctx" not in name: 
+            if "ctx" not in name and "prompt_integrator" not in name: 
                param.requires_grad_(False)
            else:
                print(name)
@@ -322,8 +386,10 @@ class MSGCoOp(TrainerX):
            load_pretrained_weights(self.model.prompt_learner, cfg.MODEL.INIT_WEIGHTS)

        self.model.to(self.device)
-        # NOTE: only give prompt_learner to the optimizer
-        self.optim = build_optimizer(self.model.prompt_learner, cfg.OPTIM)
+        
+        # NOTE: give prompt_learner and prompt_integrator to the optimizer
+        trainable_params = list(self.model.prompt_learner.parameters()) + list(self.model.prompt_integrator.parameters())
+        self.optim = build_optimizer([{'params': trainable_params}], cfg.OPTIM)
        self.sched = build_lr_scheduler(self.optim, cfg.OPTIM)
        self.register_model("prompt_learner", self.model.prompt_learner, self.optim, self.sched)
        
@@ -352,8 +418,12 @@ class MSGCoOp(TrainerX):
            self.scaler.step(self.optim)
            self.scaler.update()
        else:
-            output, score, diversity_loss = self.model(image)
-            loss = F.cross_entropy(output, label)+self.w*score + diversity_loss * self.diversity_weight
+            output, score, diversity_loss, hybrid_weights, entropy = self.model(image)
+            
+            # Add attention regularization to encourage balanced prompt usage
+            attn_reg = -(hybrid_weights * torch.log(hybrid_weights + 1e-8)).mean()
+            
+            loss = F.cross_entropy(output, label) + self.w * score + diversity_loss * self.diversity_weight + self.attn_reg_weight * attn_reg
            self.model_backward_and_update(loss)

        loss_summary = {
Author	SHA1	Message	Date
rain-bus	17c6c4c309	Uncertain fuse	2026-02-01 17:48:47 +08:00
rain-bus	0ba13ffbbd	Attn fuse	2026-01-31 23:48:05 +08:00