Skip to content
字数
1623 字
阅读时间
7 分钟

1. 项目目标

我的研究目标是训练一个下一最佳视角(Next-Best-View, NBV)策略网络 。该网络能够主动指导一个通用的三维重建模型 MapAnything,通过迭代式地预测信息最丰富的下一个相机位姿,从而以最少的视图数量,最大化最终三维重建的质量。

2. 核心组件

我的研究框架主要由以下三个部分构成:

  • 三维重建骨干 (MapAnything): 我采用现有的 MapAnything 作为一个强大、统一且前馈式的三维重建骨干。它的核心是一个为多视图几何设计的24层交替注意力Transformer

    • 输入与预处理: 模型首先使用 DINOv2 将输入的 张图像编码为图像块令牌(Patch Tokens)。同时,它能够灵活地融合为任意视图子集提供的可选几何信息,如相机位姿、深度图、内参。一个特殊的可学习尺度令牌会被附加到这 组视图令牌之后,同时,一个固定的参考视图嵌入会被加到第一个视图的令牌上,以区分参考坐标系。

    • 核心融合机制: Transformer 通过其交替注意力的结构,在视图内(Intra-View)和视图间(Inter-View)两个维度上高效地融合所有 个视图的信息。

    • 分解式输出: Transformer 的输出令牌会被送入不同的解码头,分别预测一个全局度量尺度因子 (),以及每个视图局部的光线方向 ()、up-to-scale 的光线深度 () 和相机位姿 ()。最终的度量点云是通过将这些分解式的输出组合计算而成的。

  • NBV 策略网络 : 这是我需要训练的核心模块。它的输入,即场景特征 ,正是从 MapAnything 的多视图Transformer 的最后一层的输出令牌。该特征张量的形状为 ,其中 是批次大小, 是初始视图数, 是每个视图的图像块数量,具体来说是518/14*518/14=1369,而 则是 Transformer 的隐层维度。我的网络接收这个高维特征 ,在维度上平均池化后,回归出下一个最佳相机位姿(相对第一个View的相对位姿) (目前我只预测三维坐标,旋转自动朝向物体)。

    • 代码如下:
python
class AttentionNBVPolicy(BaseNBVPolicy):
    """
    基于注意力机制的NBV策略网络
    
    使用Transformer编码器处理序列特征
    """
    
    def __init__(self, 
                 scene_feature_dim: int = 768,
                 hidden_dim: int = 512,
                 num_heads: int = 8,
                 num_layers: int = 4,
                 output_mode: str = "cartesian",
                 token_pooling_mode: str = "mean"):
        """
        初始化注意力NBV策略网络
        
        Args:
            scene_feature_dim: 场景特征维度
            hidden_dim: 隐藏层维度
            num_heads: 注意力头数
            num_layers: Transformer层数
            output_mode: 输出模式
        """
        super().__init__(output_mode, token_pooling_mode)
        
        self.scene_feature_dim = scene_feature_dim
        self.hidden_dim = hidden_dim
        
        # 特征投影
        self.feature_projection = nn.Linear(scene_feature_dim, hidden_dim)
        
        # 位置编码
        self.pos_embedding = nn.Parameter(torch.randn(1, 100, hidden_dim))  # 支持最多100个视角
        
        # Transformer编码器
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=hidden_dim,
            nhead=num_heads,
            dim_feedforward=hidden_dim * 4,
            dropout=0.1,
            batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)
        
        # 全局特征提取
        self.global_pool = nn.MultiheadAttention(
            embed_dim=hidden_dim,
            num_heads=num_heads,
            batch_first=True
        )
        self.global_token = nn.Parameter(torch.randn(1, 1, hidden_dim))
        
        # 输出头
        self.output_head = nn.Sequential(
            nn.LayerNorm(hidden_dim),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim // 2, self.target_dim)
        )
        
        self._initialize_weights()
    
    def forward(self, scene_features: torch.Tensor) -> torch.Tensor:
        """
        前向传播
        
        Args:
            scene_features: 场景特征 [B, S, 768]
            
        Returns:
            camera_pose: 相机位姿 [B, target_dim]
        """
        scene_features = self._pool_tokens_if_needed(scene_features)  # [B, S, D]
        B, S, D = scene_features.shape
        
        # 特征投影
        x = self.feature_projection(scene_features)  # [B, S, hidden_dim]
        
        # 添加位置编码
        x = x + self.pos_embedding[:, :S, :]
        
        # Transformer编码
        encoded_features = self.transformer(x)  # [B, S, hidden_dim]
        
        # 全局特征提取
        global_token = self.global_token.expand(B, -1, -1)
        global_features, _ = self.global_pool(global_token, encoded_features, encoded_features)
        global_features = global_features.squeeze(1)  # [B, hidden_dim]
        
        # 输出预测
        nbv_raw = self.output_head(global_features)
        
        return self._activate_nbv(nbv_raw)
  • 可微分渲染器 : 我基于 PyTorch3D 构建了一个可微分渲染器,用于根据给定的真值三维网格 和一个相机位姿 ,合成一个新的观测视图

3. 训练流程

我采用端到端监督学习的方式来训练我的策略网络。

我最初的方案试图让梯度同时通过可微分渲染器和 MapAnything 模型本身,但发现来自渲染器的梯度路径极不稳定,会产生数值爆炸。我通过将渲染器输出的 V_{N+1} 从计算图中分离 (.detach())——切断了这条不稳定的路径。

我目前的训练流程如下:

  1. 状态表征: 给定 个初始视图(是RGB图、深度图和内参) ,我将它们送入 MapAnything 提取场景特征

  2. 策略预测: 我的策略网络 接收 并回归出下一个相机位姿:

  3. 环境交互与梯度阻断: 使用 合成新视图,并阻断其梯度:

  4. 重建与评估: 我将完整的视图集 和位姿集 一同送入 MapAnything,得到最终的预测点云

  5. 损失函数: 使用倒角距离(Chamfer Distance, CD) 作为最终的损失函数,它直接衡量了预测点云 与真值点云 之间的几何相似度。其定义如下:

    这个损失函数为我的策略网络提供了一个明确的、端到端的优化目标,即生成的位姿 必须能引导 MapAnything 产生一个在几何上更接近真值的重建结果。

  6. 反向传播:

4. 当前状态与挑战

  • 已验证: 在“单 mesh + 固定1张初始照片”的设定下,我的模型 loss 能够稳定下降。这证明了我的“模型内部梯度路径”是通畅且可学习的,验证了整个 pipeline 的正确性。

  • 当前挑战: 当我切换到“单 mesh + 随机1张初始照片”的设定时,loss 下降变得不稳定。

贡献者

页面历史