dformerv2

dformerv2

DFormerv2: Geometry-aware RGB+Depth encoder for imitation learning.

Based on: "DFormerv2: Geometry Self-Attention for RGBD Semantic Segmentation" https://github.com/VCIP-RGBD/DFormer

DFormerVariant

Bases: StrEnum

Available DFormerv2 model variants.

DFormerPretrainedWeights

Bases: StrEnum

Pretrained checkpoint families on the HuggingFace mirror.

DFormerStage

DFormerStage(embedding_dimension, number_of_heads, num_blocks, decomposition_mode, drop_path_rate=0.0, use_layer_scale=False, layer_scale_init_value=1e-05, initial_decay=2.0, decay_range=4.0, ffn_expansion_factor=4, downsample=None, use_raster_positions=False, use_feedforward_convolution=False)

Bases: Module

Single DFormer stage with multiple geometric attention blocks and optional downsampling.

Initialize DFormer stage.

Parameters:

Name	Type	Description	Default
embedding_dimension	int	Feature dimension for this stage	required
number_of_heads	int	Number of attention heads	required
num_blocks	int	Number of geometric attention blocks in this stage	required
decomposition_mode	AttentionDecompositionMode	Attention computation strategy (full or separable)	required
drop_path_rate	float	Stochastic depth rate	0.0
use_layer_scale	bool	Whether to use layer scaling	False
layer_scale_init_value	float	Initial value for layer scale parameters	1e-05
initial_decay	float	Initial decay rate for spatial biases	2.0
decay_range	float	Range of decay rates across heads	4.0
ffn_expansion_factor	int	Expansion factor for FFN hidden dimension	4
downsample	Module | None	Optional downsampling module for next stage	None
use_raster_positions	bool	Whether rotary encoding uses flattened raster grid positions (the DFormerv2 reference convention).	False
use_feedforward_convolution	bool	Whether blocks use the DFormerv2 FFN with an inner depthwise convolution instead of a plain MLP.	False

Source code in src/versatil/models/encoding/encoders/cross_modal/rgbd/dformerv2.py

def __init__(
    self,
    embedding_dimension: int,
    number_of_heads: int,
    num_blocks: int,
    decomposition_mode: AttentionDecompositionMode,
    drop_path_rate: float = 0.0,
    use_layer_scale: bool = False,
    layer_scale_init_value: float = 1e-5,
    initial_decay: float = 2.0,
    decay_range: float = 4.0,
    ffn_expansion_factor: int = 4,
    downsample: nn.Module | None = None,
    use_raster_positions: bool = False,
    use_feedforward_convolution: bool = False,
):
    """Initialize DFormer stage.

    Args:
        embedding_dimension: Feature dimension for this stage
        number_of_heads: Number of attention heads
        num_blocks: Number of geometric attention blocks in this stage
        decomposition_mode: Attention computation strategy (full or separable)
        drop_path_rate: Stochastic depth rate
        use_layer_scale: Whether to use layer scaling
        layer_scale_init_value: Initial value for layer scale parameters
        initial_decay: Initial decay rate for spatial biases
        decay_range: Range of decay rates across heads
        ffn_expansion_factor: Expansion factor for FFN hidden dimension
        downsample: Optional downsampling module for next stage
        use_raster_positions: Whether rotary encoding uses flattened raster
            grid positions (the DFormerv2 reference convention).
        use_feedforward_convolution: Whether blocks use the DFormerv2 FFN
            with an inner depthwise convolution instead of a plain MLP.
    """
    super().__init__()
    self.embedding_dimension = embedding_dimension

    self.blocks = nn.ModuleList(
        [
            GeometricAttentionEncoderBlock(
                decomposition_mode=decomposition_mode,
                embedding_dimension=embedding_dimension,
                number_of_heads=number_of_heads,
                ffn_dimension=embedding_dimension * ffn_expansion_factor,
                drop_path_rate=drop_path_rate,
                use_layer_scale=use_layer_scale,
                layer_scale_init_value=layer_scale_init_value,
                initial_decay=initial_decay,
                decay_range=decay_range,
                use_raster_positions=use_raster_positions,
                use_feedforward_convolution=use_feedforward_convolution,
            )
            for _ in range(num_blocks)
        ]
    )

    self.downsample = downsample
    self.norm = nn.LayerNorm(embedding_dimension, eps=1e-6)

forward

forward(rgb_features, depth_map)

Forward pass through the stage.

Parameters:

Name	Type	Description	Default
rgb_features	Tensor	RGB features of shape (B, H, W, C)	required
depth_map	Tensor	Depth map of shape (B, 1, H, W)	required

Returns:

Name	Type	Description
output_features	Tensor	Normalized features for this stage output (B, H, W, C)
next_features	Tensor	Downsampled features for next stage (B, H', W', C') or same as output if no downsample
depth_map	Tensor	Resized depth map for next stage (B, 1, H', W') or same as input if no downsample

Source code in src/versatil/models/encoding/encoders/cross_modal/rgbd/dformerv2.py

def forward(
    self, rgb_features: torch.Tensor, depth_map: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Forward pass through the stage.

    Args:
        rgb_features: RGB features of shape (B, H, W, C)
        depth_map: Depth map of shape (B, 1, H, W)

    Returns:
        output_features: Normalized features for this stage output (B, H, W, C)
        next_features: Downsampled features for next stage (B, H', W', C') or same as output if no downsample
        depth_map: Resized depth map for next stage (B, 1, H', W') or same as input if no downsample
    """
    features = rgb_features
    for block in self.blocks:
        features = block(features, depth_map)
    output_features = self.norm(features)
    if self.downsample is not None:
        next_features = self.downsample(features)
    else:
        next_features = output_features
    return output_features, next_features, depth_map

DFormerEncoder

DFormerEncoder(input_keys, variant=value, decomposition_mode=value, drop_path_rate=0.1, layer_scale_init_value=1e-06, initial_decay=2.0, pretrained=False, frozen=False, pretrained_weights=value, pooling_method=value, model_dtype=None, lora_config=None)

Bases: RGBDEncoderMixin, Encoder

DFormerv2 encoder for RGB+Depth fusion using geometric self-attention.

Hierarchical encoder with multi-scale feature extraction and depth-conditioned attention.

Initialize DFormer encoder.

Parameters:

Name	Type	Description	Default
input_keys	str | list[str]	Input keys for RGB and depth	required
variant	str	Model variant (S/B/L)	value
decomposition_mode	str	Attention computation strategy	value
pooling_method	str	Feature pooling method (spatial_softmax or global_average)	value
drop_path_rate	float	Stochastic depth rate	0.1
layer_scale_init_value	float	Initial value for layer scale	1e-06
initial_decay	float	Initial decay rate for spatial biases	2.0
pretrained	bool	Whether to use pretrained weights	False
frozen	bool	Whether to freeze encoder weights	False
pretrained_weights	str	Which checkpoint family to download from https://huggingface.co/bbynku/DFormerv2 when pretrained is set: the ImageNet backbone or the NYU/SUNRGBD finetuned models.	value
model_dtype	str | None	Precision string from experiment config (e.g. "bf16-mixed").	None
lora_config	LoRAAdaptation | None	Optional LoRA adapter configuration applied to the stage linears.	None

Source code in src/versatil/models/encoding/encoders/cross_modal/rgbd/dformerv2.py

def __init__(
    self,
    input_keys: str | list[str],
    variant: str = DFormerVariant.SMALL.value,
    decomposition_mode: str = AttentionDecompositionMode.SEPARABLE.value,
    drop_path_rate: float = 0.1,
    layer_scale_init_value: float = 1e-6,
    initial_decay: float = 2.0,
    pretrained: bool = False,
    frozen: bool = False,
    pretrained_weights: str = DFormerPretrainedWeights.IMAGENET.value,
    pooling_method: str = PoolingMethod.AVERAGE.value,
    model_dtype: str | None = None,
    lora_config: LoRAAdaptation | None = None,
):
    """Initialize DFormer encoder.

    Args:
        input_keys: Input keys for RGB and depth
        variant: Model variant (S/B/L)
        decomposition_mode: Attention computation strategy
        pooling_method: Feature pooling method (spatial_softmax or global_average)
        drop_path_rate: Stochastic depth rate
        layer_scale_init_value: Initial value for layer scale
        initial_decay: Initial decay rate for spatial biases
        pretrained: Whether to use pretrained weights
        frozen: Whether to freeze encoder weights
        pretrained_weights: Which checkpoint family to download from
            https://huggingface.co/bbynku/DFormerv2 when ``pretrained``
            is set: the ImageNet backbone or the NYU/SUNRGBD finetuned
            models.
        model_dtype: Precision string from experiment config (e.g. ``"bf16-mixed"``).
        lora_config: Optional LoRA adapter configuration applied to the
            stage linears.
    """
    specification = EncoderInput(
        keys=input_keys,
        exactly_one_camera_modality=[CameraModality.RGB, CameraModality.DEPTH],
        required_camera_modalities=[CameraModality.RGB, CameraModality.DEPTH],
    )
    super().__init__(
        input_specification=specification,
        pretrained=pretrained,
        frozen=frozen,
        model_dtype=model_dtype,
    )
    if variant not in self.VARIANT_CONFIGS:
        raise ValueError(
            f"Variant '{variant}' not supported. "
            f"Choose from: {list(self.VARIANT_CONFIGS.keys())}"
        )
    weights_key = (variant, DFormerPretrainedWeights(pretrained_weights).value)
    self._setup_camera_keys(input_keys=self.input_specification.keys)
    self.variant = variant
    self.pooling_method = pooling_method
    self.decomposition_mode = AttentionDecompositionMode(decomposition_mode)
    config = self.VARIANT_CONFIGS[variant]
    self.embed_dims: list[int] = config["embed_dims"]
    self.depths: list[int] = config["depths"]
    self.number_of_heads: list[int] = config["number_of_heads"]
    self.decay_ranges: list[int] = config["decay_ranges"]
    self.ffn_ratios: list[int] = config["ffn_ratios"]
    self.use_layer_scales: list[bool] = config["use_layer_scales"]
    self.num_stages = len(self.embed_dims)
    # Patch size is fixed at 4 to match original DFormerv2 architecture (2 stride-2 convs = 4x downsample)
    # This cannot be changed without breaking pretrained model loading
    self.patch_embed = PatchEmbedding(
        patch_size=4,
        in_chans=3,
        embed_dim=self.embed_dims[0],
        embed_type=PatchEmbedType.PROGRESSIVE.value,
        norm_layer=FrozenBatchNorm2d,
    )
    self._build_backbone(
        drop_path_rate=drop_path_rate,
        layer_scale_init_value=layer_scale_init_value,
        initial_decay=initial_decay,
    )
    self.feature_dim = self.embed_dims[-1]
    self.pooling_head: PoolingHead | None = None
    self.output_dim: int | tuple[int, ...] = self.feature_dim
    if pretrained:
        checkpoint_path = hf_hub_download(
            repo_id=DFORMER_HUGGINGFACE_REPO,
            filename=DFORMER_PRETRAINED_FILENAMES[weights_key],
        )
        self._load_checkpoint(checkpoint_path)
    self.lora_config = lora_config
    self.stages = nn.ModuleList(
        [
            apply_lora_config(model=stage, lora_config=lora_config, frozen=frozen)
            for stage in self.stages
        ]
    )
    if is_lora_enabled(lora_config=lora_config):
        # PEFT freezes the wrapped stages; the patch embedding sits outside
        # them and must freeze too so only adapters train.
        for parameter in self.patch_embed.parameters():
            parameter.requires_grad = False
    if frozen:
        super()._freeze_weights()
    self._apply_model_dtype()

encode

encode(inputs)

Encode RGB + depth through DFormer.

Parameters:

Name	Type	Description	Default
inputs	dict[str, Tensor]	Dict with RGB as (B, C, H, W) and depth as (B, 1, H, W).	required

Returns:

Type	Description
dict[str, Tensor]	Dict with RGBD features.

Source code in src/versatil/models/encoding/encoders/cross_modal/rgbd/dformerv2.py

def encode(self, inputs: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
    """Encode RGB + depth through DFormer.

    Args:
        inputs: Dict with RGB as (B, C, H, W) and depth as (B, 1, H, W).

    Returns:
        Dict with RGBD features.
    """
    rgb_key = self._camera_key_for_modality(modality=CameraModality.RGB)
    depth_key = self._camera_key_for_modality(modality=CameraModality.DEPTH)

    rgb = inputs[rgb_key]
    depth_map = inputs[depth_key]
    rgb_features = self.patch_embed(rgb)  # (B, H_patches, W_patches, C)
    features = rgb_features
    for stage in self.stages:
        output_features, next_features, depth_map = stage(features, depth_map)
        features = next_features

    if self.pooling_head is None:
        raise RuntimeError(
            "pooling_head is not initialized. Call set_image_size() before forward."
        )
    final_features = features.permute(0, 3, 1, 2).contiguous()
    pooled_features = self.pooling_head(final_features)
    return {EncoderOutputKeys.RGBD.value: pooled_features}

set_image_size

set_image_size(image_height, image_width)

Compute feature map dimensions and create pooling head.

Parameters:

Name	Type	Description	Default
image_height	int	Target image height.	required
image_width	int	Target image width.	required

Source code in src/versatil/models/encoding/encoders/cross_modal/rgbd/dformerv2.py

def set_image_size(self, image_height: int, image_width: int) -> None:
    """Compute feature map dimensions and create pooling head.

    Args:
        image_height: Target image height.
        image_width: Target image width.
    """
    mock_input_dtype = self._mock_forward_dtype()
    with torch.no_grad(), self._mock_forward_autocast():
        mock_rgb = torch.zeros(
            1, 3, image_height, image_width, dtype=mock_input_dtype
        )
        features = self.patch_embed(mock_rgb)
        depth_map = torch.zeros(
            1, 1, image_height, image_width, dtype=mock_input_dtype
        )
        for stage in self.stages:
            _, features, depth_map = stage(features, depth_map)
        _, spatial_height, spatial_width, _ = features.shape
    self._setup_pooling(spatial_height=spatial_height, spatial_width=spatial_width)
    if self.frozen:
        self._freeze_weights()
    self._apply_model_dtype()

get_explainability_targets

get_explainability_targets()

Return the final DFormer stage for spatial attribution maps.

DFormer stages return a tuple; output_index=0 selects the stage output before downsampling. The output is NHWC and is converted by the explainability package before feature-grid map computation.

Returns:

Type	Description
list[VisionExplanationTarget]	One NHWC spatial feature-map target for the final DFormer stage.

Source code in src/versatil/models/encoding/encoders/cross_modal/rgbd/dformerv2.py

def get_explainability_targets(self) -> list[VisionExplanationTarget]:
    """Return the final DFormer stage for spatial attribution maps.

    DFormer stages return a tuple; ``output_index=0`` selects the stage
    output before downsampling. The output is NHWC and is converted by the
    explainability package before feature-grid map computation.

    Returns:
        One NHWC spatial feature-map target for the final DFormer stage.
    """
    return [
        VisionExplanationTarget(
            layer=self.stages[-1],
            target_kind=ExplanationTargetKind.SPATIAL_FEATURE_MAP.value,
            activation_layout=ActivationLayout.NHWC.value,
            output_index=0,
        )
    ]

get_output_specification

get_output_specification()

Return the output feature names and dimensions for this encoder.

Returns:

Type	Description
list[FeatureMetadata]	List of FeatureMetadata with RGBD feature name and its pooled dimension.

Source code in src/versatil/models/encoding/encoders/cross_modal/rgbd/dformerv2.py

def get_output_specification(self) -> list[FeatureMetadata]:
    """Return the output feature names and dimensions for this encoder.

    Returns:
        List of FeatureMetadata with RGBD feature name and its pooled dimension.
    """
    dimension = (
        (self.output_dim,) if isinstance(self.output_dim, int) else self.output_dim
    )
    return [
        FeatureMetadata(
            key=EncoderOutputKeys.RGBD.value,
            feature_type=infer_feature_type(dimension),
            dimension=dimension,
        )
    ]