Logic
467 lines
20 KiB
Python
467 lines
20 KiB
Python
from roboimi.vla.core.interfaces import VLABackbone
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
import torchvision
|
||
import numpy as np
|
||
from typing import Callable, Optional, Tuple, Union
|
||
|
||
from .attnres_resnet2d import AttnResResNetLikeBackbone2D
|
||
|
||
def _replace_submodules(
|
||
root_module: nn.Module, predicate: Callable[[nn.Module], bool], func: Callable[[nn.Module], nn.Module]
|
||
) -> nn.Module:
|
||
"""
|
||
Args:
|
||
root_module: 需要替换子模块的根模块
|
||
predicate: 接受一个模块作为参数，如果该模块需要被替换则返回 True。
|
||
func: 接受一个模块作为参数，并返回一个新的模块来替换它。
|
||
Returns:
|
||
子模块已被替换的根模块。
|
||
"""
|
||
if predicate(root_module):
|
||
return func(root_module)
|
||
|
||
replace_list = [k.split(".") for k, m in root_module.named_modules(remove_duplicate=True) if predicate(m)]
|
||
for *parents, k in replace_list:
|
||
parent_module = root_module
|
||
if len(parents) > 0:
|
||
parent_module = root_module.get_submodule(".".join(parents))
|
||
if isinstance(parent_module, nn.Sequential):
|
||
src_module = parent_module[int(k)]
|
||
else:
|
||
src_module = getattr(parent_module, k)
|
||
tgt_module = func(src_module)
|
||
if isinstance(parent_module, nn.Sequential):
|
||
parent_module[int(k)] = tgt_module
|
||
else:
|
||
setattr(parent_module, k, tgt_module)
|
||
# 验证所有 BN 是否已被替换
|
||
assert not any(predicate(m) for _, m in root_module.named_modules(remove_duplicate=True))
|
||
return root_module
|
||
|
||
class SpatialSoftmax(nn.Module):
|
||
"""
|
||
Finn 等人在 "Deep Spatial Autoencoders for Visuomotor Learning" 中描述的空间软 Argmax 操作
|
||
(https://huggingface.co/papers/1509.06113)。这是 robomimic 实现的一个最小移植版本。
|
||
"""
|
||
|
||
def __init__(self, input_shape, num_kp=None):
|
||
"""
|
||
Args:
|
||
input_shape (list): (C, H, W) 输入特征图形状。
|
||
num_kp (int): 输出中的关键点数量。如果为 None，输出将具有与输入相同的通道数。
|
||
"""
|
||
super().__init__()
|
||
|
||
assert len(input_shape) == 3
|
||
self._in_c, self._in_h, self._in_w = input_shape
|
||
|
||
if num_kp is not None:
|
||
self.nets = torch.nn.Conv2d(self._in_c, num_kp, kernel_size=1)
|
||
self._out_c = num_kp
|
||
else:
|
||
self.nets = None
|
||
self._out_c = self._in_c
|
||
|
||
# 我们可以直接使用 torch.linspace，但这似乎与 numpy 的行为略有不同
|
||
# 并且会导致预训练模型的 pc_success 略有下降。
|
||
pos_x, pos_y = np.meshgrid(np.linspace(-1.0, 1.0, self._in_w), np.linspace(-1.0, 1.0, self._in_h))
|
||
pos_x = torch.from_numpy(pos_x.reshape(self._in_h * self._in_w, 1)).float()
|
||
pos_y = torch.from_numpy(pos_y.reshape(self._in_h * self._in_w, 1)).float()
|
||
# 注册为 buffer，以便将其移动到正确的设备。
|
||
self.register_buffer("pos_grid", torch.cat([pos_x, pos_y], dim=1))
|
||
|
||
def forward(self, features: torch.Tensor) -> torch.Tensor:
|
||
"""
|
||
Args:
|
||
features: (B, C, H, W) 输入特征图。
|
||
Returns:
|
||
(B, K, 2) 关键点的图像空间坐标。
|
||
"""
|
||
if self.nets is not None:
|
||
features = self.nets(features)
|
||
|
||
# [B, K, H, W] -> [B * K, H * W]，其中 K 是关键点数量
|
||
features = features.reshape(-1, self._in_h * self._in_w)
|
||
# 2d softmax 归一化
|
||
attention = F.softmax(features, dim=-1)
|
||
# [B * K, H * W] x [H * W, 2] -> [B * K, 2] 用于 x 和 y 维度的空间坐标均值
|
||
expected_xy = attention @ self.pos_grid
|
||
# 重塑为 [B, K, 2]
|
||
feature_keypoints = expected_xy.view(-1, self._out_c, 2)
|
||
|
||
return feature_keypoints
|
||
|
||
class _SingleRgbEncoder(nn.Module):
|
||
"""单个摄像头的 RGB 编码器，支持独立或共享使用"""
|
||
def __init__(
|
||
self,
|
||
vision_backbone: str,
|
||
pretrained_backbone_weights: str | None,
|
||
input_shape: Tuple[int, int, int],
|
||
crop_shape: Optional[Tuple[int, int]],
|
||
crop_is_random: bool,
|
||
use_group_norm: bool,
|
||
spatial_softmax_num_keypoints: int,
|
||
freeze_backbone: bool = True, # 新增：是否冻结backbone
|
||
vision_backbone_mode: str = "resnet",
|
||
attnres_stem_dim: int = 64,
|
||
attnres_stage_dims: Optional[Tuple[int, ...]] = None,
|
||
attnres_stage_depths: Optional[Tuple[int, ...]] = None,
|
||
attnres_stage_heads: Optional[Tuple[int, ...]] = None,
|
||
attnres_stage_kv_heads: Optional[Tuple[int, ...]] = None,
|
||
attnres_stage_window_sizes: Optional[Tuple[int, ...]] = None,
|
||
attnres_dropout: float = 0.0,
|
||
attnres_ffn_mult: float = 2.667,
|
||
attnres_eps: float = 1e-6,
|
||
attnres_rope_theta: float = 10000.0,
|
||
):
|
||
super().__init__()
|
||
|
||
# 设置可选的预处理
|
||
if crop_shape is not None:
|
||
self.do_crop = True
|
||
# 评估时始终使用中心裁剪
|
||
self.center_crop = torchvision.transforms.CenterCrop(crop_shape)
|
||
if crop_is_random:
|
||
self.maybe_random_crop = torchvision.transforms.RandomCrop(crop_shape)
|
||
else:
|
||
self.maybe_random_crop = self.center_crop
|
||
else:
|
||
self.do_crop = False
|
||
crop_shape = input_shape[1:]
|
||
|
||
if vision_backbone_mode == "resnet":
|
||
# 设置骨干网络
|
||
backbone_model = getattr(torchvision.models, vision_backbone)(
|
||
weights=pretrained_backbone_weights
|
||
)
|
||
|
||
# 移除 AvgPool 和 FC (假设 layer4 是 children()[-3])
|
||
self.backbone = nn.Sequential(*(list(backbone_model.children())[:-2]))
|
||
|
||
if use_group_norm:
|
||
self.backbone = _replace_submodules(
|
||
root_module=self.backbone,
|
||
predicate=lambda x: isinstance(x, nn.BatchNorm2d),
|
||
func=lambda x: nn.GroupNorm(
|
||
num_groups=max(1, x.num_features // 16),
|
||
num_channels=x.num_features,
|
||
),
|
||
)
|
||
elif vision_backbone_mode == "attnres_resnet":
|
||
self.backbone = AttnResResNetLikeBackbone2D(
|
||
input_channels=input_shape[0],
|
||
stem_dim=attnres_stem_dim,
|
||
stage_dims=tuple(attnres_stage_dims or (64, 128, 256, 512)),
|
||
stage_depths=tuple(attnres_stage_depths or (2, 2, 2, 2)),
|
||
stage_heads=tuple(attnres_stage_heads or (4, 4, 8, 8)),
|
||
stage_kv_heads=tuple(attnres_stage_kv_heads or (1, 1, 1, 1)),
|
||
stage_window_sizes=tuple(attnres_stage_window_sizes or (7, 7, 7, 7)),
|
||
use_group_norm=use_group_norm,
|
||
dropout=attnres_dropout,
|
||
ffn_mult=attnres_ffn_mult,
|
||
eps=attnres_eps,
|
||
rope_theta=attnres_rope_theta,
|
||
)
|
||
else:
|
||
raise ValueError(f"不支持的 vision_backbone_mode: {vision_backbone_mode}")
|
||
|
||
# 冻结backbone参数（可选）
|
||
if freeze_backbone:
|
||
for param in self.backbone.parameters():
|
||
param.requires_grad = False
|
||
|
||
# 设置池化和最终层
|
||
# 使用试运行来获取特征图形状
|
||
dummy_shape = (1, input_shape[0], *crop_shape)
|
||
with torch.no_grad():
|
||
dummy_out = self.backbone(torch.zeros(dummy_shape))
|
||
feature_map_shape = dummy_out.shape[1:] # (C, H, W)
|
||
|
||
self.pool = SpatialSoftmax(feature_map_shape, num_kp=spatial_softmax_num_keypoints)
|
||
self.feature_dim = spatial_softmax_num_keypoints * 2
|
||
self.out = nn.Linear(spatial_softmax_num_keypoints * 2, self.feature_dim)
|
||
self.relu = nn.ReLU()
|
||
|
||
# 注册ImageNet标准化参数为buffer（会自动移到GPU）
|
||
self.register_buffer('mean', torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
|
||
self.register_buffer('std', torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
|
||
|
||
def forward_single_image(self, x: torch.Tensor) -> torch.Tensor:
|
||
if self.do_crop:
|
||
x = self.maybe_random_crop(x) if self.training else self.center_crop(x)
|
||
|
||
# ImageNet标准化（预训练权重期望的输入分布）
|
||
x = (x - self.mean) / self.std
|
||
|
||
x = self.relu(self.out(torch.flatten(self.pool(self.backbone(x)), start_dim=1)))
|
||
return x
|
||
|
||
|
||
class ResNetDiffusionBackbone(VLABackbone):
|
||
def __init__(
|
||
self,
|
||
vision_backbone: str = "resnet18",
|
||
pretrained_backbone_weights: str | None = None,
|
||
input_shape: Tuple[int, int, int] = (3, 84, 84), # (C, H, W)
|
||
crop_shape: Optional[Tuple[int, int]] = None,
|
||
crop_is_random: bool = True,
|
||
use_group_norm: bool = True,
|
||
spatial_softmax_num_keypoints: int = 32,
|
||
use_separate_rgb_encoder_per_camera: bool = False, # 新增：是否为每个摄像头使用独立编码器
|
||
output_tokens_per_camera: bool = False, # 是否按相机返回多token，而不是拼成一个token
|
||
num_cameras: int = 1, # 新增：摄像头数量（仅在独立编码器模式下使用）
|
||
camera_names: Optional[Tuple[str, ...]] = None, # 显式相机顺序
|
||
freeze_backbone: bool = True, # 新增：是否冻结ResNet backbone（推荐True）
|
||
vision_backbone_mode: str = "resnet",
|
||
attnres_stem_dim: int = 64,
|
||
attnres_stage_dims: Optional[Tuple[int, ...]] = None,
|
||
attnres_stage_depths: Optional[Tuple[int, ...]] = None,
|
||
attnres_stage_heads: Optional[Tuple[int, ...]] = None,
|
||
attnres_stage_kv_heads: Optional[Tuple[int, ...]] = None,
|
||
attnres_stage_window_sizes: Optional[Tuple[int, ...]] = None,
|
||
attnres_dropout: float = 0.0,
|
||
attnres_ffn_mult: float = 2.667,
|
||
attnres_eps: float = 1e-6,
|
||
attnres_rope_theta: float = 10000.0,
|
||
):
|
||
super().__init__()
|
||
|
||
self.use_separate_rgb_encoder_per_camera = use_separate_rgb_encoder_per_camera
|
||
self.output_tokens_per_camera = bool(output_tokens_per_camera)
|
||
self.num_cameras = num_cameras
|
||
self.tokens_per_step = self.num_cameras if self.output_tokens_per_camera else 1
|
||
self.camera_names = tuple(camera_names) if camera_names is not None else None
|
||
if self.camera_names is not None and len(self.camera_names) != self.num_cameras:
|
||
raise ValueError(
|
||
f"camera_names 长度({len(self.camera_names)})与 num_cameras({self.num_cameras})不一致"
|
||
)
|
||
|
||
if use_separate_rgb_encoder_per_camera:
|
||
# 独立编码器模式：为每个摄像头创建独立的编码器
|
||
encoders = [
|
||
_SingleRgbEncoder(
|
||
vision_backbone=vision_backbone,
|
||
pretrained_backbone_weights=pretrained_backbone_weights,
|
||
input_shape=input_shape,
|
||
crop_shape=crop_shape,
|
||
crop_is_random=crop_is_random,
|
||
use_group_norm=use_group_norm,
|
||
spatial_softmax_num_keypoints=spatial_softmax_num_keypoints,
|
||
freeze_backbone=freeze_backbone,
|
||
vision_backbone_mode=vision_backbone_mode,
|
||
attnres_stem_dim=attnres_stem_dim,
|
||
attnres_stage_dims=attnres_stage_dims,
|
||
attnres_stage_depths=attnres_stage_depths,
|
||
attnres_stage_heads=attnres_stage_heads,
|
||
attnres_stage_kv_heads=attnres_stage_kv_heads,
|
||
attnres_stage_window_sizes=attnres_stage_window_sizes,
|
||
attnres_dropout=attnres_dropout,
|
||
attnres_ffn_mult=attnres_ffn_mult,
|
||
attnres_eps=attnres_eps,
|
||
attnres_rope_theta=attnres_rope_theta,
|
||
)
|
||
for _ in range(num_cameras)
|
||
]
|
||
self.rgb_encoder = nn.ModuleList(encoders)
|
||
# 重要：output_dim 始终表示单个编码器的特征维度（与 lerobot 保持一致）
|
||
self.feature_dim = encoders[0].feature_dim
|
||
else:
|
||
# 共享编码器模式：所有摄像头共享同一个编码器
|
||
self.rgb_encoder = _SingleRgbEncoder(
|
||
vision_backbone=vision_backbone,
|
||
pretrained_backbone_weights=pretrained_backbone_weights,
|
||
input_shape=input_shape,
|
||
crop_shape=crop_shape,
|
||
crop_is_random=crop_is_random,
|
||
use_group_norm=use_group_norm,
|
||
spatial_softmax_num_keypoints=spatial_softmax_num_keypoints,
|
||
freeze_backbone=freeze_backbone,
|
||
vision_backbone_mode=vision_backbone_mode,
|
||
attnres_stem_dim=attnres_stem_dim,
|
||
attnres_stage_dims=attnres_stage_dims,
|
||
attnres_stage_depths=attnres_stage_depths,
|
||
attnres_stage_heads=attnres_stage_heads,
|
||
attnres_stage_kv_heads=attnres_stage_kv_heads,
|
||
attnres_stage_window_sizes=attnres_stage_window_sizes,
|
||
attnres_dropout=attnres_dropout,
|
||
attnres_ffn_mult=attnres_ffn_mult,
|
||
attnres_eps=attnres_eps,
|
||
attnres_rope_theta=attnres_rope_theta,
|
||
)
|
||
self.feature_dim = self.rgb_encoder.feature_dim
|
||
|
||
def _ordered_camera_names(self, images) -> Tuple[str, ...]:
|
||
if self.camera_names is None:
|
||
camera_names = tuple(sorted(images.keys()))
|
||
if len(camera_names) != self.num_cameras:
|
||
raise ValueError(
|
||
f"图像输入相机数量({len(camera_names)})与 num_cameras({self.num_cameras})不一致"
|
||
)
|
||
return camera_names
|
||
|
||
missing = [cam_name for cam_name in self.camera_names if cam_name not in images]
|
||
if missing:
|
||
raise ValueError(
|
||
f"图像输入缺少必需相机。missing={missing}, expected={list(self.camera_names)}"
|
||
)
|
||
return self.camera_names
|
||
|
||
def forward(self, images):
|
||
"""
|
||
Args:
|
||
images: Dict[str, Tensor], 每个摄像头的图像
|
||
形状: {cam_name: (B, T, C, H, W)}
|
||
|
||
Returns:
|
||
Tensor: (B, T, total_feature_dim)
|
||
"""
|
||
any_tensor = next(iter(images.values()))
|
||
B, T = any_tensor.shape[:2]
|
||
cam_names = self._ordered_camera_names(images)
|
||
|
||
features_all = []
|
||
if self.use_separate_rgb_encoder_per_camera:
|
||
# 独立编码器模式：每个摄像头使用对应的编码器
|
||
for cam_idx, cam_name in enumerate(cam_names):
|
||
img = images[cam_name]
|
||
encoder = self.rgb_encoder[cam_idx]
|
||
features = encoder.forward_single_image(img.reshape(B * T, *img.shape[2:]))
|
||
features_all.append(features)
|
||
else:
|
||
# 共享编码器模式：所有摄像头共享同一个编码器
|
||
for cam_name in cam_names:
|
||
img = images[cam_name]
|
||
features = self.rgb_encoder.forward_single_image(img.reshape(B * T, *img.shape[2:]))
|
||
features_all.append(features)
|
||
|
||
if self.output_tokens_per_camera:
|
||
stacked = torch.stack(features_all, dim=1) # (B*T, num_cams, feature_dim)
|
||
return stacked.view(B, T, len(cam_names), self.feature_dim)
|
||
return torch.cat(features_all, dim=1).view(B, T, -1)
|
||
|
||
@property
|
||
def output_dim(self):
|
||
return self.feature_dim
|
||
|
||
if __name__ == "__main__":
|
||
print("=" * 60)
|
||
print("🚀 Testing ResNetDiffusionBackbone")
|
||
print("=" * 60)
|
||
|
||
# Configuration
|
||
B, T = 2, 5
|
||
C, H, W = 3, 96, 96
|
||
crop_h, crop_w = 84, 84
|
||
num_keypoints = 32
|
||
feature_dim_per_cam = num_keypoints * 2
|
||
|
||
# Create dummy input (2 cameras)
|
||
images = {
|
||
"cam_high": torch.randn(B, T, C, H, W),
|
||
"cam_wrist": torch.randn(B, T, C, H, W)
|
||
}
|
||
num_cameras = len(images)
|
||
|
||
# ============================================================================
|
||
# Test 1: Shared Encoder (默认模式)
|
||
# ============================================================================
|
||
print("\n[Test 1] Shared Encoder Mode")
|
||
print("-" * 60)
|
||
backbone_shared = ResNetDiffusionBackbone(
|
||
vision_backbone="resnet18",
|
||
pretrained_backbone_weights=None, # Speed up test
|
||
input_shape=(C, H, W),
|
||
crop_shape=(crop_h, crop_w),
|
||
crop_is_random=True,
|
||
use_group_norm=True,
|
||
spatial_softmax_num_keypoints=num_keypoints,
|
||
use_separate_rgb_encoder_per_camera=False, # 共享编码器
|
||
)
|
||
|
||
print(f"✅ Shared encoder model instantiated")
|
||
print(f" Output dim per camera: {feature_dim_per_cam}")
|
||
print(f" Number of cameras: {num_cameras}")
|
||
print(f" Expected total dim: {num_cameras * feature_dim_per_cam}")
|
||
|
||
output = backbone_shared(images)
|
||
print(f"\n🔄 Forward pass completed")
|
||
print(f" Input shapes: {[v.shape for v in images.values()]}")
|
||
print(f" Output shape: {output.shape}")
|
||
|
||
expected_dim = num_cameras * feature_dim_per_cam
|
||
assert output.shape == (B, T, expected_dim), f"Expected shape {(B, T, expected_dim)}, got {output.shape}"
|
||
print(f"✨ Test passed!")
|
||
|
||
# ============================================================================
|
||
# Test 2: Separate Encoders (独立编码器模式)
|
||
# ============================================================================
|
||
print("\n[Test 2] Separate Encoders Mode")
|
||
print("-" * 60)
|
||
backbone_separate = ResNetDiffusionBackbone(
|
||
vision_backbone="resnet18",
|
||
pretrained_backbone_weights=None, # Speed up test
|
||
input_shape=(C, H, W),
|
||
crop_shape=(crop_h, crop_w),
|
||
crop_is_random=True,
|
||
use_group_norm=True,
|
||
spatial_softmax_num_keypoints=num_keypoints,
|
||
use_separate_rgb_encoder_per_camera=True, # 独立编码器
|
||
num_cameras=num_cameras,
|
||
)
|
||
|
||
print(f"✅ Separate encoders model instantiated")
|
||
print(f" Output dim per camera: {feature_dim_per_cam}")
|
||
print(f" Number of cameras: {num_cameras}")
|
||
print(f" Number of encoders: {len(backbone_separate.rgb_encoder)}")
|
||
|
||
output = backbone_separate(images)
|
||
print(f"\n🔄 Forward pass completed")
|
||
print(f" Input shapes: {[v.shape for v in images.values()]}")
|
||
print(f" Output shape: {output.shape}")
|
||
|
||
expected_dim = num_cameras * feature_dim_per_cam
|
||
assert output.shape == (B, T, expected_dim), f"Expected shape {(B, T, expected_dim)}, got {output.shape}"
|
||
print(f"✨ Test passed!")
|
||
|
||
# ============================================================================
|
||
# Test 3: Verify parameters count
|
||
# ============================================================================
|
||
print("\n[Test 3] Parameter Count Comparison")
|
||
print("-" * 60)
|
||
shared_params = sum(p.numel() for p in backbone_shared.parameters())
|
||
separate_params = sum(p.numel() for p in backbone_separate.parameters())
|
||
|
||
print(f" Shared encoder parameters: {shared_params:,}")
|
||
print(f" Separate encoders parameters: {separate_params:,}")
|
||
print(f" Ratio: {separate_params / shared_params:.2f}x")
|
||
|
||
assert separate_params > shared_params, "Separate encoders should have more parameters"
|
||
print(f"✨ Verification passed!")
|
||
|
||
# ============================================================================
|
||
# Test 4: Verify independent parameters
|
||
# ============================================================================
|
||
print("\n[Test 4] Verify Independent Parameters")
|
||
print("-" * 60)
|
||
# Check that encoders have independent parameters
|
||
encoder_0_first_param = list(backbone_separate.rgb_encoder[0].parameters())[0]
|
||
encoder_1_first_param = list(backbone_separate.rgb_encoder[1].parameters())[0]
|
||
|
||
# Modify first encoder's parameter
|
||
with torch.no_grad():
|
||
encoder_0_first_param += 1.0
|
||
|
||
# Verify they are not the same tensor
|
||
assert not torch.allclose(encoder_0_first_param, encoder_1_first_param), \
|
||
"Encoders should have independent parameters"
|
||
|
||
print(f"✅ Encoders have independent parameters")
|
||
print(f"✨ All tests passed!")
|
||
|
||
print("\n" + "=" * 60)
|
||
print("🎉 All tests completed successfully!")
|
||
print("=" * 60)
|