SpatialRGBEncoder(input_keys, backbone=value, pooling_method=value, batch_norm_handling=value, intermediate_layer_index=None, pretrained=False, frozen=False, model_dtype=None, lora_config=None)
Bases: RGBEncoderMixin, SpatialBackboneEncoder
RGB encoder for backbones that output spatial feature maps.
Supports any timm backbone compatible with features_only=True,
regardless of whether the architecture is convolutional (ResNet,
EfficientNet, ConvNeXt) or attention-based (Swin, TinyViT).
Handles both NCHW and NHWC output layouts transparently.
Source code in src/versatil/models/encoding/encoders/spatial_backbone.py
| def __init__(
self,
input_keys: str | list[str],
backbone: str = SpatialBackboneType.RESNET18.value,
pooling_method: str = PoolingMethod.AVERAGE.value,
batch_norm_handling: str = BatchNormHandling.FROZEN.value,
intermediate_layer_index: int | None = None,
pretrained: bool = False,
frozen: bool = False,
model_dtype: str | None = None,
lora_config: LoRAAdaptation | None = None,
) -> None:
"""Initialize the spatial encoder with a timm backbone.
Args:
input_keys: Camera observation keys.
backbone: timm model name from SpatialBackboneType.
pooling_method: Feature pooling strategy.
batch_norm_handling: How to handle batch normalization layers.
intermediate_layer_index: Optional timm intermediate layer index
to pool. Negative values index from the end; ``None`` uses
the last layer.
pretrained: Whether to load pretrained weights.
frozen: Whether to freeze all parameters.
model_dtype: Precision string from experiment config (e.g. ``"bf16-mixed"``).
lora_config: Optional PEFT LoRA adapter configuration.
"""
specification = EncoderInput(
keys=input_keys,
required_camera_modalities=[self._camera_modality],
)
super().__init__(
input_specification=specification,
pretrained=pretrained,
frozen=frozen,
model_dtype=model_dtype,
)
valid_backbones = [e.value for e in SpatialBackboneType]
if backbone not in valid_backbones:
raise ValueError(
f"Invalid backbone '{backbone}'. Must be one of: {valid_backbones}"
)
pooling = PoolingMethod(pooling_method)
if not pooling.supports_spatial:
raise ValueError(
f"Pooling method '{pooling_method}' is not compatible with "
f"spatial feature maps. Use one of: "
f"{[p.value for p in PoolingMethod if p.supports_spatial]}"
)
self._setup_camera_keys(input_keys=self.input_specification.keys)
self.batch_norm_handling = batch_norm_handling
self.pooling_method = pooling_method
self.intermediate_layer_index = intermediate_layer_index
self.backbone_name = backbone
self.lora_config = lora_config
self._channels_last = False
self._build_backbone()
self.feature_dim = self._get_intermediate_layer_channels()
self.pooling_head: PoolingHead | None = None
self.output_dim: int | tuple[int, ...] = self.feature_dim
if frozen:
super()._freeze_weights()
self._apply_model_dtype()
|