核心思想是将普通的ViT编码器应用在多个尺度提取的块上,并将块预测融合到一个端到端的可训练模型中的单个高分辨率密集预测中。架构如图3所示。
为了预测深度,作者使用了两个ViT编码器,一个块编码器和一个图像编码器。
-
块编码器应用于多个尺度下提取的块。从直接上将,这可能允许学习的表示是尺度不变的,因为权重是跨尺度共享的。
-
图像编码器将块预测锚定在全局上下文中。他应用于整个输入图像, 下采样到所选择的编码器主干的基本输入分辨率(本文是384×384)。
整个网络的固定分辨率为 1536×1536,他被选定为 ViT 的 384×384 的倍数。这就保证了足够大的感受野和对任意图像的恒定运行时间,同时防止了超出显存问题。
降采样到 1536×1536,768×768 和 384 ×384 后,输入图像被分割成大小为 384×384 的块,为了避免接缝,对于最精细的两个尺度,作者分割时采用重叠划分的方式。在每个尺度下,块被输入到块编码器中,为每个块(图3中的特征3-5)产生一个 24 ×24 大小的特征张量,在最精细的尺度上作者进一步提取中间特征(图3中的特征1和2)以捕获更细粒度的细节。
这里感觉原文说的不是很清楚,可以看代码 https://github.com/apple/ml-depth-pro/blob/main/src/depth_pro/network/encoder.py#L247
代码中的
x_latent0_features
,x_latent1_features
,x0_features
,x1_features
,x3_features
分别表示图 3 中的 feature 1-5。# Step 0: create a 3-level image pyramid. x0, x1, x2 = self._create_pyramid(x) # Step 1: split to create batched overlapped mini-images at the backbone (BeiT/ViT/Dino) resolution. # 5x5 @ 384x384 at the highest resolution (1536x1536). x0_patches = self.split(x0, overlap_ratio=0.25) # 3x3 @ 384x384 at the middle resolution (768x768). x1_patches = self.split(x1, overlap_ratio=0.5) # 1x1 # 384x384 at the lowest resolution (384x384). x2_patches = x2 # Concatenate all the sliding window patches and form a batch of size (35=5x5+3x3+1x1). x_pyramid_patches = torch.cat( (x0_patches, x1_patches, x2_patches), dim=0, ) # Step 2: Run the backbone (BeiT) model and get the result of large batch size. x_pyramid_encodings = self.patch_encoder(x_pyramid_patches) x_pyramid_encodings = self.reshape_feature( x_pyramid_encodings, self.out_size, self.out_size ) # Step 3: merging. # Merge highres latent encoding. x_latent0_encodings = self.reshape_feature( self.backbone_highres_hook0, # 中间特征 self.out_size, self.out_size, ) x_latent0_features = self.merge( x_latent0_encodings[: batch_size * 5 * 5], batch_size=batch_size, padding=3 ) x_latent1_encodings = self.reshape_feature( self.backbone_highres_hook1, # 中间特征 self.out_size, self.out_size, ) x_latent1_features = self.merge( x_latent1_encodings[: batch_size * 5 * 5], batch_size=batch_size, padding=3 ) # Split the 35 batch size from pyramid encoding back into 5x5+3x3+1x1. x0_encodings, x1_encodings, x2_encodings = torch.split( x_pyramid_encodings, [len(x0_patches), len(x1_patches), len(x2_patches)], dim=0, ) # 96x96 feature maps by merging 5x5 @ 24x24 patches with overlaps. x0_features = self.merge(x0_encodings, batch_size=batch_size, padding=3) # 48x84 feature maps by merging 3x3 @ 24x24 patches with overlaps. x1_features = self.merge(x1_encodings, batch_size=batch_size, padding=6) # 24x24 feature maps. x2_features = x2_encodings
过了Patch编码器后,共 (5×5×2) + (5×5+3×3+1) 个patch。
5×5×2 表示,作者用最精细的1536分辨率上拆分的 5×5 的块用编码器编码时,保留了中间 2 层的特征,并将其输出,按重叠划分的逆过程又合并回去
5×5+3×3+1表示,各个尺度下的特征经过编码后,最后一层输出的特征,然后再合并回去
除了跨尺度共享表示外,编码器网络的基于patch的应用支持并行化,因为patch可以独立处理。计算效率的另一个来源是,与将ViT扩展到更高分辨率相比,基于patch处理的计算复杂度更低。因为多头注意力计算复杂度是输入像素数量的二次方,因此在图像维度上呈4次方关系。