Add configurable beta_reduce_mode for TTT3R updates across all eval scripts

datasets_preprocess/long_prepare_bonn.py

@@ -4,10 +4,10 @@
 import numpy as np
 
 START_FRAME = 30  # inital frame
-for TARGET_FRAMES in [50,100,150,200,250,300,350,400,450,500]: 
+for TARGET_FRAMES in [50,100,150,200,250,300,350,400,450,500]:
     END_FRAME = START_FRAME + TARGET_FRAMES   # end frame
 
-    dirs = glob.glob("/home/xingyu/monst3r/data/bonn/rgbd_bonn_dataset/*/")
+    dirs = glob.glob("./data/bonn/rgbd_bonn_dataset/*/")
     dirs = sorted(dirs)
 
     # create new base directory

datasets_preprocess/long_prepare_kitti.py

@@ -21,7 +21,7 @@ def depth_read(filename):
 import shutil
 
 for TARGET_FRAMES in [50,100,150,200,250,300,350,400,450,500]:
-    depth_dirs = glob.glob("/home/xingyu/monst3r/data/kitti/val/*/proj_depth/groundtruth/image_02")
+    depth_dirs = glob.glob("./data/kitti/val/*/proj_depth/groundtruth/image_02")
     for dir in depth_dirs:
         # new depth dir
         new_depth_dir = f"./data/long_kitti_s1/depth_selection/val_selection_cropped/groundtruth_depth_gathered_{TARGET_FRAMES}/" + dir.split("/")[-4]+"_02"

demo.py

@@ -95,6 +95,13 @@ def parse_args():
         default="cut3r",
         help="model update type: cut3r or ttt3r",
     )
+    parser.add_argument(
+        "--beta_reduce_mode",
+        type=str,
+        default="mean",
+        choices=["mean", "max", "variance", "entropy"],
+        help="How TTT3R aggregates cross-attention into per-token update weights.",
+    )
     parser.add_argument(
         "--frame_interval",
         type=int,
@@ -463,6 +470,7 @@ def run_inference(args):
     print(f"Loading model from {args.model_path}...")
     model = ARCroco3DStereo.from_pretrained(args.model_path).to(device)
     model.config.model_update_type = args.model_update_type
+    model.config.beta_reduce_mode = args.beta_reduce_mode
 
     model.eval()
 

eval/mv_recon/launch.py

@@ -40,6 +40,9 @@ def get_args_parser():
     parser.add_argument("--freeze", action="store_true")
     parser.add_argument("--max_frames", type=int, default=None, help="max frames limit")
     parser.add_argument("--model_update_type", type=str, default="cut3r", help="model update type")
+    parser.add_argument("--beta_reduce_mode", type=str, default="mean",
+                        choices=["mean", "max", "variance", "entropy"],
+                        help="TTT3R beta aggregation mode")
     parser.add_argument("--voxel_size", type=float, default=0.0, help="voxel size for voxel grid downsampling, 0 means no downsampling")
     return parser
 
@@ -58,7 +61,7 @@ def main(args):
     datasets_all = {
         "7scenes": SevenScenes(
             split="test",
-            ROOT="/home/share/Dataset/3D_scene/7scenes/", # "./data/7scenes",
+            ROOT="./data/7scenes",
             resolution=resolution,
             num_seq=1,
             full_video=True,
@@ -103,6 +106,7 @@ def main(args):
 
     model = ARCroco3DStereo.from_pretrained(args.weights).to(device)
     model.config.model_update_type = args.model_update_type
+    model.config.beta_reduce_mode = args.beta_reduce_mode
 
     model.eval()
     # else:
@@ -183,7 +187,7 @@ def main(args):
                     with torch.cuda.amp.autocast(enabled=False):
                         start = time.time()
                         output = model(batch)
-                        # preds, batch = model.forward_recurrent_light(batch)
+                        # preds, batch = model.forward_recurrent_light(batch)  
                         end = time.time()
                         preds, batch = output.ress, output.views
                     valid_length = len(preds) // revisit

eval/mv_recon/run.sh

@@ -4,23 +4,28 @@ set -e
 
 workdir='.'
 model_names=('ttt3r') # ttt3r cut3r
+beta_reduce_modes=('mean' 'entropy') # mean max variance entropy
 ckpt_name='cut3r_512_dpt_4_64'
 model_weights="${workdir}/src/${ckpt_name}.pth"
 
 for model_name in "${model_names[@]}"; do
+for beta_reduce_mode in "${beta_reduce_modes[@]}"; do
 
 # for max_frames in 50 100 150 200 250 300 350 400
 for max_frames in 200
 
 do
-    output_dir="${workdir}/eval_results/video_recon/7scenes_${max_frames}/${model_name}"
+
+    output_dir="${workdir}/eval_results/video_recon/7scenes_${max_frames}/${model_name}_${beta_reduce_mode}"
     echo "$output_dir"
     NCCL_TIMEOUT=360000 accelerate launch --num_processes 1 --main_process_port 29502 eval/mv_recon/launch.py \
         --weights "$model_weights" \
         --output_dir "$output_dir" \
         --model_name "$model_name" \
         --model_update_type "$model_name" \
-        --max_frames $max_frames \
+        --beta_reduce_mode "$beta_reduce_mode" \
+        --max_frames $max_frames
 
 done
 done
+done

eval/relpose/launch.py

@@ -54,6 +54,14 @@ def get_args_parser():
         help="model type for state update strategy: cut3r or ttt3r",
     )
 
+    parser.add_argument(
+        "--beta_reduce_mode",
+        type=str,
+        default="mean",
+        choices=["mean", "max", "variance", "entropy"],
+        help="How TTT3R aggregates cross-attention into per-token update weights.",
+    )
+
     parser.add_argument(
         "--pose_eval_stride", default=1, type=int, help="stride for pose evaluation"
     )
@@ -461,5 +469,6 @@ def prepare_output(outputs, revisit=1, solve_pose=False):
 
     # set model type
     model.config.model_update_type = args.model_update_type
+    model.config.beta_reduce_mode = args.beta_reduce_mode
 
     eval_pose_estimation(args, model, save_dir=args.output_dir)

eval/relpose/run_tum.sh

@@ -4,6 +4,7 @@ set -e
 
 workdir='.'
 model_names=('ttt3r') # ttt3r cut3r
+beta_reduce_modes=('mean' 'entropy') # mean max variance entropy
 
 ckpt_name='cut3r_512_dpt_4_64'
 model_weights="${workdir}/src/${ckpt_name}.pth"
@@ -13,15 +14,18 @@ model_weights="${workdir}/src/${ckpt_name}.pth"
 datasets=('tum_s1_1000')
 
 for model_name in "${model_names[@]}"; do
+for beta_reduce_mode in "${beta_reduce_modes[@]}"; do
 for data in "${datasets[@]}"; do
-    output_dir="${workdir}/eval_results/relpose/${data}/${model_name}"
+    output_dir="${workdir}/eval_results/relpose/${data}/${model_name}_${beta_reduce_mode}"
     echo "$output_dir"
     accelerate launch --num_processes 2 --main_process_port 29551 eval/relpose/launch.py \
         --weights "$model_weights" \
         --output_dir "$output_dir" \
         --eval_dataset "$data" \
         --size 512 \
-        --model_update_type "$model_name"
+        --model_update_type "$model_name" \
+        --beta_reduce_mode "$beta_reduce_mode"
+done
 done
 done
 

eval/video_depth/launch.py

@@ -53,6 +53,14 @@ def get_args_parser():
         help="model type for state update strategy: cut3r or ttt3r",
     )
 
+    parser.add_argument(
+        "--beta_reduce_mode",
+        type=str,
+        default="mean",
+        choices=["mean", "max", "variance", "entropy"],
+        help="How TTT3R aggregates cross-attention into per-token update weights.",
+    )
+
 
     parser.add_argument(
         "--pose_eval_stride", default=1, type=int, help="stride for pose evaluation"
@@ -339,5 +347,6 @@ def prepare_output(outputs, revisit=1):
 
     # set model type
     model.config.model_update_type = args.model_update_type
+    model.config.beta_reduce_mode = args.beta_reduce_mode
 
     eval_pose_estimation(args, model, save_dir=args.output_dir)

eval/video_depth/run_bonn.sh

@@ -4,23 +4,26 @@ set -e
 
 workdir='.'
 model_names=('ttt3r') # ttt3r cut3r
+beta_reduce_modes=('mean' 'entropy') # mean max variance entropy
 ckpt_name='cut3r_512_dpt_4_64'
 model_weights="${workdir}/src/${ckpt_name}.pth"
 # datasets=('bonn_s1_50' 'bonn_s1_100' 'bonn_s1_110' 'bonn_s1_150' 'bonn_s1_200' 'bonn_s1_250' 'bonn_s1_300' 'bonn_s1_350' 'bonn_s1_400' 'bonn_s1_450' 'bonn_s1_500')
 datasets=('bonn_s1_500')
 
 
 for model_name in "${model_names[@]}"; do
+for beta_reduce_mode in "${beta_reduce_modes[@]}"; do
 for data in "${datasets[@]}"; do
-    output_dir="${workdir}/eval_results/video_depth/${data}/${model_name}"
+    output_dir="${workdir}/eval_results/video_depth/${data}/${model_name}_${beta_reduce_mode}"
     echo "$output_dir"
 
     accelerate launch --num_processes 1 --main_process_port 29556 eval/video_depth/launch.py \
         --weights "$model_weights" \
         --output_dir "$output_dir" \
         --eval_dataset "$data" \
         --size 512 \
-        --model_update_type "$model_name"
+        --model_update_type "$model_name" \
+        --beta_reduce_mode "$beta_reduce_mode"
 
     # scale&shift scale metric
     python eval/video_depth/eval_depth.py \
@@ -39,3 +42,4 @@ for data in "${datasets[@]}"; do
     --align "scale&shift"
 done
 done
+done

eval/video_depth/run_kitti.sh

@@ -4,23 +4,26 @@ set -e
 
 workdir='.'
 model_names=('ttt3r') # ttt3r cut3r
+beta_reduce_modes=('mean' 'entropy') # mean max variance entropy
 ckpt_name='cut3r_512_dpt_4_64'
 model_weights="${workdir}/src/${ckpt_name}.pth"
 # datasets=('kitti_s1_50' 'kitti_s1_100' 'kitti_s1_110' 'kitti_s1_150' 'kitti_s1_200' 'kitti_s1_250' 'kitti_s1_300' 'kitti_s1_350' 'kitti_s1_400' 'kitti_s1_450' 'kitti_s1_500')
 datasets=('kitti_s1_500')
 
 
 for model_name in "${model_names[@]}"; do
+for beta_reduce_mode in "${beta_reduce_modes[@]}"; do
 for data in "${datasets[@]}"; do
-    output_dir="${workdir}/eval_results/video_depth/${data}/${model_name}"
+    output_dir="${workdir}/eval_results/video_depth/${data}/${model_name}_${beta_reduce_mode}"
     echo "$output_dir"
 
     accelerate launch --num_processes 1 --main_process_port 29555 eval/video_depth/launch.py \
         --weights "$model_weights" \
         --output_dir "$output_dir" \
         --eval_dataset "$data" \
         --size 512 \
-        --model_update_type "$model_name"
+        --model_update_type "$model_name" \
+        --beta_reduce_mode "$beta_reduce_mode"
 
     # scale&shift scale metric
     python eval/video_depth/eval_depth.py \
@@ -39,3 +42,4 @@ for data in "${datasets[@]}"; do
     --align "scale&shift"
 done
 done
+done

eval/video_depth/run_sintel.sh

@@ -4,22 +4,25 @@ set -e
 
 workdir='.'
 model_names=('ttt3r') # ttt3r cut3r
+beta_reduce_modes=('mean' 'entropy') # mean max variance entropy
 
 ckpt_name='cut3r_512_dpt_4_64'
 model_weights="${workdir}/src/${ckpt_name}.pth"
 datasets=('sintel')
 
 for model_name in "${model_names[@]}"; do
+for beta_reduce_mode in "${beta_reduce_modes[@]}"; do
 for data in "${datasets[@]}"; do
-    output_dir="${workdir}/eval_results/video_depth/${data}/${model_name}"
+    output_dir="${workdir}/eval_results/video_depth/${data}/${model_name}_${beta_reduce_mode}"
     echo "$output_dir"
 
     accelerate launch --num_processes 1 --main_process_port 29555 eval/video_depth/launch.py \
         --weights "$model_weights" \
         --output_dir "$output_dir" \
         --eval_dataset "$data" \
         --size 512 \
-        --model_update_type "$model_name"
+        --model_update_type "$model_name" \
+        --beta_reduce_mode "$beta_reduce_mode"
 
     # scale&shift scale metric
     python eval/video_depth/eval_depth.py \
@@ -38,3 +41,4 @@ for data in "${datasets[@]}"; do
     --align "scale&shift"
 done
 done
+done

eval/video_depth/utils.py

@@ -107,7 +107,7 @@ def get_vertical_colorbar(h, vmin, vmax, cmap_name="jet", label=None, cbar_preci
 
     # Do some plotting.
     ax = fig.add_subplot(111)
-    cmap = cm.get_cmap(cmap_name)
+    cmap = mpl.colormaps[cmap_name]
     norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
 
     tick_cnt = 6
@@ -177,7 +177,7 @@ def colorize_np(
     x = (x - vmin) / (vmax - vmin)
     # x = np.clip(x, 0., 1.)
 
-    cmap = cm.get_cmap(cmap_name)
+    cmap = mpl.colormaps[cmap_name]
     x_new = cmap(x)[:, :, :3]
 
     if mask is not None:

src/dust3r/model.py

@@ -74,7 +74,9 @@ def strip_module(state_dict):
 def load_model(model_path, device, verbose=True):
     if verbose:
         print("... loading model from", model_path)
-    ckpt = torch.load(model_path, map_location="cpu")
+    # PyTorch 2.6+ defaults to weights_only=True, which breaks older
+    # checkpoints storing OmegaConf/config objects alongside weights.
+    ckpt = torch.load(model_path, map_location="cpu", weights_only=False)
     args = ckpt["args"].model.replace(
         "ManyAR_PatchEmbed", "PatchEmbedDust3R"
     )  # ManyAR only for aspect ratio not consistent
@@ -117,6 +119,7 @@ def __init__(
         pose_conf_head=False,
         pose_head=False,
         model_update_type="cut3r",
+        beta_reduce_mode="mean",
         **croco_kwargs,
     ):
         super().__init__()
@@ -138,6 +141,7 @@ def __init__(
         self.pose_conf_head = pose_conf_head
         self.pose_head = pose_head
         self.model_update_type = model_update_type
+        self.beta_reduce_mode = beta_reduce_mode
         self.croco_kwargs = croco_kwargs
 
 
@@ -326,6 +330,30 @@ def from_pretrained(cls, pretrained_model_name_or_path, **kw):
                 )
             return model
 
+    def _compute_ttt3r_beta(self, cross_attn_state):
+        cross_attn_state = rearrange(
+            torch.cat(cross_attn_state, dim=0),
+            "l h nstate nimg -> 1 nstate nimg (l h)",
+        )
+        reduce_mode = getattr(self.config, "beta_reduce_mode", "mean")
+
+        if reduce_mode == "mean":
+            beta = torch.sigmoid(cross_attn_state.mean(dim=(-1, -2)))
+        elif reduce_mode == "max":
+            beta = torch.sigmoid(cross_attn_state.amax(dim=(-1, -2)))
+        elif reduce_mode == "variance":
+            beta = torch.sigmoid(cross_attn_state.var(dim=(-1, -2), unbiased=False))
+        elif reduce_mode == "entropy": 
+            attn_probs = torch.softmax(cross_attn_state, dim=2)
+            entropy = -(attn_probs * attn_probs.clamp_min(1e-8).log()).sum(dim=2)
+            norm = torch.log(torch.tensor(cross_attn_state.shape[2], device=entropy.device, dtype=entropy.dtype))
+            entropy = entropy / norm.clamp_min(1e-8)
+            beta = 1.0 - entropy.mean(dim=-1)
+        else:
+            raise ValueError(f"Invalid beta reduction mode: {reduce_mode}")
+
+        return beta[..., None]
+
     def _set_patch_embed(self, img_size=224, patch_size=16, enc_embed_dim=768):
         self.patch_embed = get_patch_embed(
             self.patch_embed_cls, img_size, patch_size, enc_embed_dim, in_chans=3
@@ -901,9 +929,7 @@ def _forward_impl(self, views, ret_state=False):
                 if self.config.model_update_type == "cut3r":
                     update_mask1 = update_mask
                 elif self.config.model_update_type == "ttt3r":
-                    cross_attn_state = rearrange(torch.cat(cross_attn_state, dim=0), 'l h nstate nimg -> 1 nstate nimg (l h)') # [12, 16, 768, 1 + 576] -> [1, 768, 1 + 576, 12*16]
-                    state_query_img_key = cross_attn_state.mean(dim=(-1, -2))
-                    update_mask1 = update_mask * torch.sigmoid(state_query_img_key)[..., None] * 1.0
+                    update_mask1 = update_mask * self._compute_ttt3r_beta(cross_attn_state)
                 else:
                     raise ValueError(f"Invalid model type: {self.config.model_update_type}")
 
@@ -1271,9 +1297,7 @@ def forward_recurrent_lighter(self, views, device='cuda', ret_state=False):
                 if self.config.model_update_type == "cut3r":
                     update_mask1 = update_mask
                 elif self.config.model_update_type == "ttt3r":
-                    cross_attn_state = rearrange(torch.cat(cross_attn_state, dim=0), 'l h nstate nimg -> 1 nstate nimg (l h)') # [12, 16, 768, 1 + 576] -> [1, 768, 1 + 576, 12*16]
-                    state_query_img_key = cross_attn_state.mean(dim=(-1, -2))
-                    update_mask1 = update_mask * torch.sigmoid(state_query_img_key)[..., None] * 1.0
+                    update_mask1 = update_mask * self._compute_ttt3r_beta(cross_attn_state)
                 else:
                     raise ValueError(f"Invalid model type: {self.config.model_update_type}")