import math import inspect from contextlib import nullcontext import torch import torch.nn as nn import torch.nn.functional as F from typing import Tuple def conv_layer(in_dim, out_dim, kernel_size=1, padding=0, stride=1): return nn.Sequential( nn.Conv2d(in_dim, out_dim, kernel_size, stride, padding, bias=False), nn.BatchNorm2d(out_dim), nn.ReLU(True)) def linear_layer(in_dim, out_dim, bias=False): return nn.Sequential(nn.Linear(in_dim, out_dim, bias), nn.BatchNorm1d(out_dim), nn.ReLU(True)) def _build_cabs_mlp(depth, hidden_size, output_hidden_size): layers = [nn.Linear(hidden_size, output_hidden_size)] for _ in range(1, depth): layers.append(nn.SiLU()) layers.append(nn.Linear(output_hidden_size, output_hidden_size)) return nn.Sequential(*layers) class LayerNorm2d(nn.Module): def __init__(self, num_channels, eps=1e-6): super().__init__() self.weight = nn.Parameter(torch.ones(num_channels)) self.bias = nn.Parameter(torch.zeros(num_channels)) self.eps = eps def forward(self, x): mean = x.mean(dim=1, keepdim=True) var = (x - mean).pow(2).mean(dim=1, keepdim=True) x = (x - mean) * torch.rsqrt(var + self.eps) return x * self.weight[:, None, None] + self.bias[:, None, None] class CAbstractorBlock(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() self.shortcut = ( nn.Identity() if in_channels == out_channels else nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False) ) self.net = nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False), LayerNorm2d(out_channels), nn.SiLU(inplace=True), nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False), LayerNorm2d(out_channels), nn.SiLU(inplace=True), nn.Conv2d(out_channels, out_channels, kernel_size=1, bias=False), LayerNorm2d(out_channels), ) self.act = nn.SiLU(inplace=True) def forward(self, x): return self.act(self.net(x) + self.shortcut(x)) class CAbstractorAdapter(nn.Module): """ Local C-Abstractor adapter for v2_C. It mirrors Honeybee's C-Abstractor shape contract without importing Honeybee: 2D visual features are optionally position encoded, processed by convolutional abstractor blocks, adaptively pooled to query tokens, then projected to the language hidden size. """ def __init__( self, in_channels, hidden_channels, out_channels, depth=3, mlp_depth=2, num_query_tokens=144, num_input_tokens=None, pos_emb=True, prenorm=False, ): super().__init__() self.in_channels = in_channels self.num_query_tokens = num_query_tokens self.input_hw = None self.output_hw = None if num_query_tokens is not None and num_query_tokens > 0: query_hw = int(math.sqrt(num_query_tokens)) if query_hw * query_hw != num_query_tokens: raise ValueError("num_query_tokens must be a square number") self.output_hw = (query_hw, query_hw) self.prenorm = nn.LayerNorm(in_channels) if prenorm else None if pos_emb and num_input_tokens is not None: self.pos_emb = nn.Parameter(torch.zeros(1, num_input_tokens, in_channels)) if hasattr(nn.init, "trunc_normal_"): nn.init.trunc_normal_(self.pos_emb, mean=0.0, std=0.02) else: nn.init.normal_(self.pos_emb, mean=0.0, std=0.02) else: self.pos_emb = None if depth > 0: first_stage = self._make_stage(depth, in_channels, hidden_channels) second_stage = self._make_stage(depth, hidden_channels, hidden_channels) self.net = nn.ModuleList([first_stage, second_stage]) self.readout = _build_cabs_mlp(mlp_depth, hidden_channels, out_channels) else: self.net = nn.ModuleList() self.readout = _build_cabs_mlp(mlp_depth, in_channels, out_channels) @staticmethod def _make_stage(depth, in_channels, out_channels): blocks = [] for idx in range(depth): blocks.append( CAbstractorBlock( in_channels if idx == 0 else out_channels, out_channels, ) ) return nn.Sequential(*blocks) def _pos_emb_for_shape(self, h, w): if self.pos_emb is None: return None if self.pos_emb.size(1) == h * w: return self.pos_emb src_hw = int(math.sqrt(self.pos_emb.size(1))) if src_hw * src_hw != self.pos_emb.size(1): raise ValueError("Cannot interpolate non-square C-Abstractor positional embeddings") pos = self.pos_emb.transpose(1, 2).reshape(1, self.in_channels, src_hw, src_hw) pos = F.interpolate(pos, size=(h, w), mode="bicubic", align_corners=False) return pos.flatten(2).transpose(1, 2) def forward(self, x): b, c, h, w = x.shape tokens = x.flatten(2).transpose(1, 2) if self.prenorm is not None: tokens = self.prenorm(tokens) pos_emb = self._pos_emb_for_shape(h, w) if pos_emb is not None: tokens = tokens + pos_emb x = tokens.transpose(1, 2).reshape(b, c, h, w) if self.net: x = self.net[0](x) target_hw = self.output_hw if self.output_hw is not None else (h, w) x = F.adaptive_avg_pool2d(x, target_hw) if self.net: x = self.net[1](x) x = x.flatten(2).transpose(1, 2) x = self.readout(x) return x, (h, w) class DAbstractorAdapter(nn.Module): """ Local deformable-attention abstractor for v2_D. This follows Honeybee's D-Abstractor interface without importing Honeybee: visual features are projected to a Deformable DETR decoder memory, learned query tokens attend to the visual grid with deformable cross-attention, and the resulting query states are projected to the language hidden size. """ def __init__( self, in_channels, out_channels, d_model=512, decoder_layers=2, num_heads=8, ffn_dim=1024, dropout=0.1, attention_dropout=0.0, num_query_tokens=144, num_input_tokens=None, pos_emb=True, pooled_v_target="query", learnable_ref_points=True, disable_custom_kernels=True, force_fp32=False, clamp_reference_points=False, reference_point_eps=1e-4, output_norm=False, output_scale=1.0, output_clip=None, sanitize_output=False, ): super().__init__() try: from transformers.models.deformable_detr import DeformableDetrConfig from transformers.models.deformable_detr.modeling_deformable_detr import ( DeformableDetrDecoderLayer, ) except ImportError as exc: raise ImportError( "DAbstractorAdapter requires transformers with Deformable DETR support." ) from exc if d_model % num_heads != 0: raise ValueError("d_abs_d_model must be divisible by d_abs_num_heads") query_hw = int(math.sqrt(num_query_tokens)) if query_hw * query_hw != num_query_tokens: raise ValueError("d_abs_num_query_tokens must be a square number") if isinstance(in_channels, (list, tuple)): self.in_channels_per_level = [int(channels) for channels in in_channels] else: self.in_channels_per_level = [int(in_channels)] self.in_channels = self.in_channels_per_level[-1] self.num_feature_levels = len(self.in_channels_per_level) self.d_model = d_model self.num_query_tokens = num_query_tokens self.pooled_v_target = pooled_v_target self.force_fp32 = force_fp32 self.clamp_reference_points = clamp_reference_points self.reference_point_eps = reference_point_eps self.output_scale = output_scale self.output_clip = output_clip self.sanitize_output = sanitize_output self.config = DeformableDetrConfig( d_model=d_model, decoder_layers=decoder_layers, decoder_attention_heads=num_heads, decoder_ffn_dim=ffn_dim, dropout=dropout, attention_dropout=attention_dropout, num_feature_levels=self.num_feature_levels, num_queries=num_query_tokens, disable_custom_kernels=disable_custom_kernels, use_pretrained_backbone=False, ) self.layers = nn.ModuleList( [DeformableDetrDecoderLayer(self.config) for _ in range(decoder_layers)] ) decoder_layer_params = inspect.signature(DeformableDetrDecoderLayer.forward).parameters self._uses_spatial_shapes_list = ( "spatial_shapes_list" in decoder_layer_params ) self.input_proj = nn.ModuleList( [ nn.Linear(channels, d_model) if channels != d_model else nn.Identity() for channels in self.in_channels_per_level ] ) if pos_emb and num_input_tokens is not None: num_input_tokens_per_level = self._normalize_num_input_tokens(num_input_tokens) pos_embs = [] for tokens in num_input_tokens_per_level: pos_emb = nn.Parameter(torch.zeros(1, tokens, d_model)) if hasattr(nn.init, "trunc_normal_"): nn.init.trunc_normal_(pos_emb, mean=0.0, std=0.02) else: nn.init.normal_(pos_emb, mean=0.0, std=0.02) pos_embs.append(pos_emb) self.pos_emb = ( pos_embs[0] if self.num_feature_levels == 1 else nn.ParameterList(pos_embs) ) else: self.pos_emb = None if self.num_feature_levels > 1: self.level_emb = nn.Parameter( torch.empty(1, self.num_feature_levels, 1, d_model) ) nn.init.normal_(self.level_emb) else: self.level_emb = None if pooled_v_target == "none": self.query_position_embeddings = nn.Embedding(num_query_tokens, d_model * 2) self.downsampler = None else: self.query_position_embeddings = nn.Embedding(num_query_tokens, d_model) self.downsampler = nn.AdaptiveAvgPool2d((query_hw, query_hw)) reference_points = self._init_reference_points(num_query_tokens, num_input_tokens) if learnable_ref_points: self.reference_points = nn.Parameter(reference_points) else: self.register_buffer("reference_points", reference_points) self.output_proj = ( nn.Linear(d_model, out_channels) if d_model != out_channels else nn.Identity() ) self.output_norm = ( nn.LayerNorm(out_channels, elementwise_affine=False) if output_norm else None ) def _normalize_num_input_tokens(self, num_input_tokens): if isinstance(num_input_tokens, (list, tuple)): if len(num_input_tokens) != self.num_feature_levels: raise ValueError( "num_input_tokens must have one value per D-Abstractor feature level" ) return [int(tokens) for tokens in num_input_tokens] if self.num_feature_levels > 1: raise ValueError( "multi-scale DAbstractorAdapter requires num_input_tokens per level" ) return [int(num_input_tokens)] @staticmethod def _init_reference_points(num_query_tokens, num_input_tokens): query_hw = int(math.sqrt(num_query_tokens)) memory_hw = query_hw if isinstance(num_input_tokens, (list, tuple)): for tokens in num_input_tokens: inferred_hw = int(math.sqrt(tokens)) if inferred_hw * inferred_hw == tokens: memory_hw = max(memory_hw, inferred_hw) elif num_input_tokens is not None: inferred_hw = int(math.sqrt(num_input_tokens)) if inferred_hw * inferred_hw == num_input_tokens: memory_hw = inferred_hw ref_y = torch.linspace(0.5, memory_hw - 0.5, query_hw, dtype=torch.float32) ref_x = torch.linspace(0.5, memory_hw - 0.5, query_hw, dtype=torch.float32) ref_y, ref_x = torch.meshgrid(ref_y, ref_x, indexing="ij") ref_y = ref_y.reshape(-1) / memory_hw ref_x = ref_x.reshape(-1) / memory_hw return torch.stack((ref_x, ref_y), dim=-1).unsqueeze(0) def _pos_emb_for_shape(self, h, w, level=0): if self.pos_emb is None: return None pos_emb = self.pos_emb[level] if isinstance(self.pos_emb, nn.ParameterList) else self.pos_emb if pos_emb.size(1) == h * w: return pos_emb src_hw = int(math.sqrt(pos_emb.size(1))) if src_hw * src_hw != pos_emb.size(1): raise ValueError("Cannot interpolate non-square D-Abstractor positional embeddings") pos = pos_emb.transpose(1, 2).reshape(1, self.d_model, src_hw, src_hw) pos = F.interpolate(pos, size=(h, w), mode="bicubic", align_corners=False) return pos.flatten(2).transpose(1, 2) @staticmethod def _prepare_deformable_inputs(batch_size, spatial_hw, dtype, device): valid_ratios = torch.ones(batch_size, len(spatial_hw), 2, dtype=dtype, device=device) spatial_shapes = torch.tensor(spatial_hw, dtype=torch.long, device=device) level_lengths = spatial_shapes[:, 0] * spatial_shapes[:, 1] level_start_index = torch.cat( [ torch.zeros(1, dtype=torch.long, device=device), level_lengths.cumsum(0)[:-1], ], dim=0, ) return valid_ratios, spatial_shapes, level_start_index def _make_pooled_queries(self, visual_tokens, h, w): batch_size = visual_tokens.size(0) query_embeds = self.query_position_embeddings.weight if self.pooled_v_target == "none": query_embed, target = torch.split(query_embeds, self.d_model, dim=1) query_embed = query_embed.unsqueeze(0).expand(batch_size, -1, -1) target = target.unsqueeze(0).expand(batch_size, -1, -1) return query_embed, target pooled = visual_tokens.transpose(1, 2).reshape(batch_size, self.d_model, h, w) pooled = self.downsampler(pooled).flatten(2).transpose(1, 2) learned_queries = query_embeds.unsqueeze(0).expand(batch_size, -1, -1) if self.pooled_v_target == "tgt": return learned_queries, pooled return pooled, learned_queries def forward(self, x): feature_maps = list(x) if isinstance(x, (list, tuple)) else [x] if len(feature_maps) != self.num_feature_levels: raise ValueError( f"Expected {self.num_feature_levels} D-Abstractor feature levels, " f"got {len(feature_maps)}" ) original_dtype = feature_maps[-1].dtype autocast_context = ( torch.cuda.amp.autocast(enabled=False) if self.force_fp32 and feature_maps[-1].is_cuda else nullcontext() ) with autocast_context: batch_size = feature_maps[0].size(0) visual_levels = [] spatial_hw = [] query_source = None query_h, query_w = None, None for level, feature_map in enumerate(feature_maps): _, _, h, w = feature_map.shape spatial_hw.append((int(h), int(w))) visual_tokens = feature_map.flatten(2).transpose(1, 2) if self.force_fp32: visual_tokens = visual_tokens.float() visual_tokens = self.input_proj[level](visual_tokens) pos_emb = self._pos_emb_for_shape(h, w, level) if pos_emb is not None: visual_tokens = visual_tokens + pos_emb.to( dtype=visual_tokens.dtype, device=visual_tokens.device ) if self.level_emb is not None: visual_tokens = visual_tokens + self.level_emb[:, level].to( dtype=visual_tokens.dtype, device=visual_tokens.device ) visual_levels.append(visual_tokens) query_source = visual_tokens query_h, query_w = h, w visual_tokens = torch.cat(visual_levels, dim=1) query_embed, target = self._make_pooled_queries(query_source, query_h, query_w) reference_points = self.reference_points.to( dtype=visual_tokens.dtype, device=visual_tokens.device ) if self.clamp_reference_points: reference_points = reference_points.clamp( self.reference_point_eps, 1.0 - self.reference_point_eps, ) reference_points = reference_points.expand(batch_size, -1, -1) valid_ratios, spatial_shapes, level_start_index = self._prepare_deformable_inputs( batch_size, spatial_hw, visual_tokens.dtype, visual_tokens.device, ) spatial_shapes_list = spatial_hw hidden_states = target for decoder_layer in self.layers: reference_points_input = reference_points[:, :, None] * valid_ratios[:, None] layer_kwargs = { "position_embeddings": query_embed, "encoder_hidden_states": visual_tokens, "reference_points": reference_points_input, "spatial_shapes": spatial_shapes, "level_start_index": level_start_index, "encoder_attention_mask": None, "output_attentions": False, } if self._uses_spatial_shapes_list: layer_kwargs["spatial_shapes_list"] = spatial_shapes_list layer_outputs = decoder_layer(hidden_states, **layer_kwargs) hidden_states = layer_outputs[0] hidden_states = self.output_proj(hidden_states) if self.output_norm is not None: hidden_states = self.output_norm(hidden_states) hidden_states = hidden_states * self.output_scale if self.output_clip is not None: hidden_states = hidden_states.clamp(-self.output_clip, self.output_clip) if self.sanitize_output: hidden_states = torch.nan_to_num( hidden_states, nan=0.0, posinf=self.output_clip if self.output_clip is not None else 1.0, neginf=-(self.output_clip if self.output_clip is not None else 1.0), ) return hidden_states.to(dtype=original_dtype), (query_h, query_w) class CoordConv(nn.Module): def __init__(self, in_channels, out_channels, kernel_size=3, padding=1, stride=1): super().__init__() self.conv1 = conv_layer(in_channels + 2, out_channels, kernel_size, padding, stride) def add_coord(self, input): b, _, h, w = input.size() x_range = torch.linspace(-1, 1, w, device=input.device) y_range = torch.linspace(-1, 1, h, device=input.device) y, x = torch.meshgrid(y_range, x_range) y = y.expand([b, 1, -1, -1]) x = x.expand([b, 1, -1, -1]) coord_feat = torch.cat([x, y], 1) input = torch.cat([input, coord_feat], 1) return input def forward(self, x): x = self.add_coord(x) x = self.conv1(x) return x class MultiTaskProjector(nn.Module): def __init__(self, word_dim=1024, in_dim=256, kernel_size=3): super().__init__() self.in_dim = in_dim self.kernel_size = kernel_size # visual projector # 视觉特征的reshape 最终输出的是 (B, 256*5, 104, 104)，表示 5 个任务的 feature map（5×256） self.vis = nn.Sequential( # os16 -> os4 nn.Upsample(scale_factor=2, mode='bilinear'), # (26x26) → (52x52) conv_layer(in_dim * 2, in_dim * 2, 3, padding=1), # 512 → 512 nn.Upsample(scale_factor=2, mode='bilinear'), # (52x52) → (104x104) conv_layer(in_dim * 2, in_dim, 3, padding=1), # 512 → 256 nn.Conv2d(in_dim, in_dim*5, 1)) # 256 → 256*5 # textual projector out_dim = 1 * in_dim * kernel_size * kernel_size + 1 self.txt = nn.Linear(word_dim, out_dim) def forward(self, x, word): ''' x: b, 512, 26, 26 word: b, 512 ''' x = self.vis(x) x = torch.tensor_split(x, 5, dim=1) # no tensor_split api in torch 1.7, please use it in higher version # x = torch.chunk(x, 5, dim=1) mask_x = x[0] grasp_qua_x = x[1] grasp_sin_x = x[2] grasp_cos_x = x[3] grasp_wid_x = x[4] B, C, H, W = mask_x.size() # 1, b*256, 104, 104 mask_x = mask_x.reshape(1, B * C, H, W) grasp_qua_x = grasp_qua_x.reshape(1, B * C, H, W) grasp_sin_x = grasp_sin_x.reshape(1, B * C, H, W) grasp_cos_x = grasp_cos_x.reshape(1, B * C, H, W) grasp_wid_x = grasp_wid_x.reshape(1, B * C, H, W) # txt: b, (256*3*3 + 1) -> b, 256, 3, 3 / b word = self.txt(word) weight, bias = word[:, :-1], word[:, -1] weight = weight.reshape(B, C, self.kernel_size, self.kernel_size) # Conv2d - 1, b*256, 104, 104 -> 1, b, 104, 104 mask_out = F.conv2d(mask_x, weight, padding=self.kernel_size // 2, groups=weight.size(0), bias=bias) grasp_qua_out = F.conv2d(grasp_qua_x, weight, padding=self.kernel_size // 2, groups=weight.size(0), bias=bias) grasp_sin_out = F.conv2d(grasp_sin_x, weight, padding=self.kernel_size // 2, groups=weight.size(0), bias=bias) grasp_cos_out = F.conv2d(grasp_cos_x, weight, padding=self.kernel_size // 2, groups=weight.size(0), bias=bias) grasp_wid_out = F.conv2d(grasp_wid_x, weight, padding=self.kernel_size // 2, groups=weight.size(0), bias=bias) mask_out = mask_out.transpose(0, 1) grasp_qua_out = grasp_qua_out.transpose(0, 1) grasp_sin_out = grasp_sin_out.transpose(0, 1) grasp_cos_out = grasp_cos_out.transpose(0, 1) grasp_wid_out = grasp_wid_out.transpose(0, 1) # b, 1, 104, 104 return mask_out, grasp_qua_out, grasp_sin_out, grasp_cos_out, grasp_wid_out class Projector(nn.Module): def __init__(self, word_dim=1024, in_dim=256, kernel_size=3): super().__init__() self.in_dim = in_dim self.kernel_size = kernel_size # visual projector self.vis = nn.Sequential( # os16 -> os4 nn.Upsample(scale_factor=2, mode='bilinear'), # 26→52 conv_layer(in_dim * 2, in_dim * 2, 3, padding=1), # 卷积层 → 提取局部模式 nn.Upsample(scale_factor=2, mode='bilinear'), # 52→104 conv_layer(in_dim * 2, in_dim, 3, padding=1), # 通道压缩 + 特征压缩 nn.Conv2d(in_dim, in_dim, 1)) # 最后一个线性变换 # textual projector # 将文本向量 word: [B, 1024] 投影为卷积核权重和偏置 输出 B x (256×3×3 + 1) 的向量 out_dim = 1 * in_dim * kernel_size * kernel_size + 1 self.txt = nn.Linear(word_dim, out_dim) def forward(self, x, word): ''' x: b, 512, 26, 26 word: b, 512 ''' x = self.vis(x) # → shape: [B, 256, 104, 104] B, C, H, W = x.size() # 1, b*256, 104, 104 x = x.reshape(1, B * C, H, W) # → [1, B×256, 104, 104] 是为了支持后面 groups=B 的逐样本动态卷积 # txt: b, (256*3*3 + 1) -> b, 256, 3, 3 / b word = self.txt(word) # [B, 2305] weight, bias = word[:, :-1], word[:, -1] weight = weight.reshape(B, C, self.kernel_size, self.kernel_size) # [B, 256, 3, 3] # Conv2d - 1, b*256, 104, 104 -> 1, b, 104, 104 out = F.conv2d(x, weight, padding=self.kernel_size // 2, groups=weight.size(0), bias=bias) out = out.transpose(0, 1) # b, 1, 104, 104 return out class TransformerDecoder(nn.Module): def __init__(self, num_layers, d_model, nhead, dim_ffn, dropout, return_intermediate=False): super().__init__() self.layers = nn.ModuleList([ TransformerDecoderLayer(d_model=d_model, nhead=nhead, dim_feedforward=dim_ffn, dropout=dropout) for _ in range(num_layers) ]) self.num_layers = num_layers self.norm = nn.LayerNorm(d_model) # 是否返回中间层输出 self.return_intermediate = return_intermediate @staticmethod def pos1d(d_model, length): """ :param d_model: dimension of the model :param length: length of positions :return: length*d_model position matrix 位置编码 """ if d_model % 2 != 0: raise ValueError("Cannot use sin/cos positional encoding with " "odd dim (got dim={:d})".format(d_model)) pe = torch.zeros(length, d_model) position = torch.arange(0, length).unsqueeze(1) div_term = torch.exp((torch.arange(0, d_model, 2, dtype=torch.float) * -(math.log(10000.0) / d_model))) pe[:, 0::2] = torch.sin(position.float() * div_term) pe[:, 1::2] = torch.cos(position.float() * div_term) return pe.unsqueeze(1) # n, 1, 512 @staticmethod def pos2d(d_model, height, width): """ :param d_model: dimension of the model :param height: height of the positions :param width: width of the positions :return: d_model*height*width position matrix """ if d_model % 4 != 0: raise ValueError("Cannot use sin/cos positional encoding with " "odd dimension (got dim={:d})".format(d_model)) pe = torch.zeros(d_model, height, width) # Each dimension use half of d_model d_model = int(d_model / 2) div_term = torch.exp( torch.arange(0., d_model, 2) * -(math.log(10000.0) / d_model)) pos_w = torch.arange(0., width).unsqueeze(1) pos_h = torch.arange(0., height).unsqueeze(1) pe[0:d_model:2, :, :] = torch.sin(pos_w * div_term).transpose( 0, 1).unsqueeze(1).repeat(1, height, 1) pe[1:d_model:2, :, :] = torch.cos(pos_w * div_term).transpose( 0, 1).unsqueeze(1).repeat(1, height, 1) pe[d_model::2, :, :] = torch.sin(pos_h * div_term).transpose( 0, 1).unsqueeze(2).repeat(1, 1, width) pe[d_model + 1::2, :, :] = torch.cos(pos_h * div_term).transpose( 0, 1).unsqueeze(2).repeat(1, 1, width) return pe.reshape(-1, 1, height * width).permute(2, 1, 0) # hw, 1, 512 def forward(self, vis, txt, pad_mask): ''' vis: b, 512, h, w txt: b, L, 512 pad_mask: b, L ''' B, C, H, W = vis.size() _, L, D = txt.size() # position encoding vis_pos = self.pos2d(C, H, W) txt_pos = self.pos1d(D, L) # reshape & permute vis = vis.reshape(B, C, -1).permute(2, 0, 1) txt = txt.permute(1, 0, 2) # forward output = vis intermediate = [] for layer in self.layers: output = layer(output, txt, vis_pos, txt_pos, pad_mask) if self.return_intermediate: # HW, b, 512 -> b, 512, HW intermediate.append(self.norm(output).permute(1, 2, 0)) if self.norm is not None: # HW, b, 512 -> b, 512, HW output = self.norm(output).permute(1, 2, 0) if self.return_intermediate: intermediate.pop() intermediate.append(output) # [output1, output2, ..., output_n] return intermediate else: # b, 512, HW return output return output class TransformerDecoderLayer(nn.Module): def __init__(self, d_model=512, nhead=9, dim_feedforward=2048, dropout=0.1): super().__init__() # Normalization Layer self.self_attn_norm = nn.LayerNorm(d_model) self.cross_attn_norm = nn.LayerNorm(d_model) # Attention Layer self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, kdim=d_model, vdim=d_model) # FFN self.ffn = nn.Sequential(nn.Linear(d_model, dim_feedforward), nn.ReLU(True), nn.Dropout(dropout), nn.LayerNorm(dim_feedforward), nn.Linear(dim_feedforward, d_model)) # LayerNorm & Dropout self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.norm3 = nn.LayerNorm(d_model) self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout) self.dropout3 = nn.Dropout(dropout) def with_pos_embed(self, tensor, pos): return tensor if pos is None else tensor + pos.to(tensor.device) def forward(self, vis, txt, vis_pos, txt_pos, pad_mask): ''' vis: 26*26, b, 512 txt: L, b, 512 vis_pos: 26*26, 1, 512 txt_pos: L, 1, 512 pad_mask: b, L ''' # Self-Attention vis2 = self.norm1(vis) q = k = self.with_pos_embed(vis2, vis_pos) vis2 = self.self_attn(q, k, value=vis2)[0] vis2 = self.self_attn_norm(vis2) vis = vis + self.dropout1(vis2) # Cross-Attention vis2 = self.norm2(vis) vis2 = self.multihead_attn(query=self.with_pos_embed(vis2, vis_pos), key=self.with_pos_embed(txt, txt_pos), value=txt, key_padding_mask=pad_mask)[0] vis2 = self.cross_attn_norm(vis2) vis = vis + self.dropout2(vis2) # FFN vis2 = self.norm3(vis) vis2 = self.ffn(vis2) vis = vis + self.dropout3(vis2) return vis class FPN(nn.Module): def __init__(self, in_channels=[512, 1024, 1024], out_channels=[256, 512, 1024]): super(FPN, self).__init__() # text projection 添加一个可学习的投影层 self.txt_proj = linear_layer(in_channels[2], out_channels[2]) # fusion 1: v5 & seq -> f_5: b, 1024, 13, 13 self.f1_v_proj = conv_layer(in_channels[2], out_channels[2], 1, 0) self.norm_layer = nn.Sequential(nn.BatchNorm2d(out_channels[2]), nn.ReLU(True)) # fusion 2: v4 & fm -> f_4: b, 512, 26, 26 self.f2_v_proj = conv_layer(in_channels[1], out_channels[1], 3, 1) self.f2_cat = conv_layer(out_channels[2] + out_channels[1], out_channels[1], 1, 0) # fusion 3: v3 & fm_mid -> f_3: b, 512, 52, 52 self.f3_v_proj = conv_layer(in_channels[0], out_channels[0], 3, 1) self.f3_cat = conv_layer(out_channels[0] + out_channels[1], out_channels[1], 1, 0) # fusion 4: f_3 & f_4 & f_5 -> fq: b, 256, 26, 26 self.f4_proj5 = conv_layer(out_channels[2], out_channels[1], 3, 1) self.f4_proj4 = conv_layer(out_channels[1], out_channels[1], 3, 1) self.f4_proj3 = conv_layer(out_channels[1], out_channels[1], 3, 1) # aggregation self.aggr = conv_layer(3 * out_channels[1], out_channels[1], 1, 0) self.coordconv = nn.Sequential( CoordConv(out_channels[1], out_channels[1], 3, 1), conv_layer(out_channels[1], out_channels[1], 3, 1)) def forward(self, imgs, state): # v3, v4, v5: 256, 52, 52 / 512, 26, 26 / 1024, 13, 13 v3, v4, v5 = imgs # fusion 1: b, 1024, 13, 13 # text projection: [B, 1024, 1, 1] state = self.txt_proj(state).unsqueeze(-1).unsqueeze( -1) # Fusion 1（v5 和文本 state 融合）： f5 = self.f1_v_proj(v5) # Conv1x1 → [B, 1024, 13, 13] f5 = self.norm_layer(f5 * state) # 广播乘文本，增强目标相关性 # fusion 2: v4 + 上采样的 f5 f4 = self.f2_v_proj(v4) # Conv3x3 → [B, 512, 26, 26] f5_ = F.interpolate(f5, scale_factor=2, mode='bilinear') # 上采样 f4 = self.f2_cat(torch.cat([f4, f5_], dim=1)) # concat + Conv1x1 → [B, 512, 26, 26] # fusion 3: v3 + avg pool 的 f4 f3 = self.f3_v_proj(v3) # Conv3x3 → [B, 256, 52, 52] f3 = F.avg_pool2d(f3, 2, 2) # 降采样 → [B, 256, 26, 26] f3 = self.f3_cat(torch.cat([f3, f4], dim=1)) # concat + Conv1x1 → [B, 512, 26, 26] # fusion 4: 融合 f3/f4/f5 fq5 = self.f4_proj5(f5) # → [B, 512, 13, 13] fq4 = self.f4_proj4(f4) # → [B, 512, 26, 26] fq3 = self.f4_proj3(f3) # → [B, 512, 26, 26] fq5 = F.interpolate(fq5, scale_factor=2, mode='bilinear') # query fq = torch.cat([fq3, fq4, fq5], dim=1) # → [B, 1536, 26, 26] fq = self.aggr(fq) # Conv1x1 聚合通道 → [B, 512, 26, 26] fq = self.coordconv(fq) # CoordConv + Conv3x3 → [B, 512, 26, 26] # b, 512, 26, 26 return fq class GeoPriorGen(nn.Module): def __init__(self, embed_dim, num_heads, initial_value, heads_range): super().__init__() angle = 1.0 / (10000 ** torch.linspace(0, 1, embed_dim // num_heads // 2)) angle = angle.unsqueeze(-1).repeat(1, 2).flatten() self.weight = nn.Parameter(torch.ones(2, 1, 1, 1), requires_grad=True) decay = torch.log( 1 - 2 ** (-initial_value - heads_range * torch.arange(num_heads, dtype=torch.float) / num_heads) ) self.register_buffer("angle", angle) self.register_buffer("decay", decay) def generate_depth_decay(self, H: int, W: int, depth_grid): """ generate 2d decay mask, the result is (HW)*(HW) H, W are the numbers of patches at each column and row """ B, _, H, W = depth_grid.shape grid_d = depth_grid.reshape(B, H * W, 1) mask_d = grid_d[:, :, None, :] - grid_d[:, None, :, :] mask_d = (mask_d.abs()).sum(dim=-1) mask_d = mask_d.unsqueeze(1) * self.decay[None, :, None, None] return mask_d def generate_pos_decay(self, H: int, W: int): """ generate 2d decay mask, the result is (HW)*(HW) H, W are the numbers of patches at each column and row """ index_h = torch.arange(H).to(self.decay) index_w = torch.arange(W).to(self.decay) grid = torch.meshgrid([index_h, index_w]) grid = torch.stack(grid, dim=-1).reshape(H * W, 2) mask = grid[:, None, :] - grid[None, :, :] mask = (mask.abs()).sum(dim=-1) mask = mask * self.decay[:, None, None] return mask def generate_1d_depth_decay(self, H, W, depth_grid): """ generate 1d depth decay mask, the result is l*l """ mask = depth_grid[:, :, :, :, None] - depth_grid[:, :, :, None, :] mask = mask.abs() mask = mask * self.decay[:, None, None, None] assert mask.shape[2:] == (W, H, H) return mask def generate_1d_decay(self, l: int): """ generate 1d decay mask, the result is l*l """ index = torch.arange(l).to(self.decay) mask = index[:, None] - index[None, :] mask = mask.abs() mask = mask * self.decay[:, None, None] return mask def forward(self, HW_tuple: Tuple[int], depth_map, split_or_not=False): """ depth_map: depth patches HW_tuple: (H, W) H * W == l """ depth_map = F.interpolate(depth_map, size=HW_tuple, mode="bilinear", align_corners=False) if split_or_not: index = torch.arange(HW_tuple[0] * HW_tuple[1]).to(self.decay) sin = torch.sin(index[:, None] * self.angle[None, :]) sin = sin.reshape(HW_tuple[0], HW_tuple[1], -1) cos = torch.cos(index[:, None] * self.angle[None, :]) cos = cos.reshape(HW_tuple[0], HW_tuple[1], -1) mask_d_h = self.generate_1d_depth_decay(HW_tuple[0], HW_tuple[1], depth_map.transpose(-2, -1)) mask_d_w = self.generate_1d_depth_decay(HW_tuple[1], HW_tuple[0], depth_map) mask_h = self.generate_1d_decay(HW_tuple[0]) mask_w = self.generate_1d_decay(HW_tuple[1]) mask_h = self.weight[0] * mask_h.unsqueeze(0).unsqueeze(2) + self.weight[1] * mask_d_h mask_w = self.weight[0] * mask_w.unsqueeze(0).unsqueeze(2) + self.weight[1] * mask_d_w geo_prior = ((sin, cos), (mask_h, mask_w)) else: index = torch.arange(HW_tuple[0] * HW_tuple[1]).to(self.decay) sin = torch.sin(index[:, None] * self.angle[None, :]) sin = sin.reshape(HW_tuple[0], HW_tuple[1], -1) cos = torch.cos(index[:, None] * self.angle[None, :]) cos = cos.reshape(HW_tuple[0], HW_tuple[1], -1) mask = self.generate_pos_decay(HW_tuple[0], HW_tuple[1]) mask_d = self.generate_depth_decay(HW_tuple[0], HW_tuple[1], depth_map) mask = self.weight[0] * mask + self.weight[1] * mask_d geo_prior = ((sin, cos), mask) return geo_prior class DWConv2d(nn.Module): def __init__(self, dim, kernel_size, stride, padding): super().__init__() self.dwconv = nn.Conv2d(dim, dim, kernel_size, stride, padding, groups=dim) def forward(self, x: torch.Tensor): """ input (b h w c) """ x = x.permute(0, 3, 1, 2) x = self.dwconv(x) x = x.permute(0, 2, 3, 1) return x def angle_transform(x, sin, cos): x1 = x[:, :, :, :, ::2] x2 = x[:, :, :, :, 1::2] #print(x1.shape, x2.shape, sin.shape, cos.shape) return (x * cos) + (torch.stack([-x2, x1], dim=-1).flatten(-2) * sin) class CrossGSA(nn.Module): def __init__(self, embed_dim, num_heads, value_factor=1): super().__init__() self.factor = value_factor self.embed_dim = embed_dim self.num_heads = num_heads self.head_dim = self.embed_dim * self.factor // num_heads self.key_dim = self.embed_dim // num_heads self.scaling = self.key_dim**-0.5 self.q_proj = nn.Linear(embed_dim, embed_dim) self.k_proj = nn.Linear(embed_dim, embed_dim) self.v_proj = nn.Linear(embed_dim, embed_dim * self.factor) self.lepe = DWConv2d(embed_dim, 5, 1, 2) self.out_proj = nn.Linear(embed_dim * self.factor, embed_dim) self.layer_norm = nn.LayerNorm(self.embed_dim, eps=1e-6) self.reset_parameters() def forward(self, x: torch.Tensor, y: torch.Tensor=None, rel_pos=None): """ x: image feature [b h w c] y: depth feature [b h w c] rel_pos: (sin, cos), mask from GeoPriorGen """ bsz, h, w, _ = x.size() q = self.q_proj(x) if y: k = self.k_proj(y) v = self.v_proj(y) else: k = self.k_proj(x) v = self.v_proj(x) lepe = self.lepe(v) k = k * self.scaling qr = q.view(bsz, h, w, self.num_heads, -1).permute(0, 3, 1, 2, 4) kr = k.view(bsz, h, w, self.num_heads, -1).permute(0, 3, 1, 2, 4) vr = v.reshape(bsz, h, w, self.num_heads, -1).permute(0, 3, 1, 2, 4) # Apply angle transform if provided if rel_pos is not None: sin, cos = rel_pos[0] qr = angle_transform(qr, sin, cos) kr = angle_transform(kr, sin, cos) qr = qr.flatten(2, 3) kr = kr.flatten(2, 3) vr = vr.flatten(2, 3) qk_mat = qr @ kr.transpose(-1, -2) if rel_pos is not None: _, mask = rel_pos qk_mat = qk_mat + mask qk_mat = torch.softmax(qk_mat, -1) output = torch.matmul(qk_mat, vr) output = output.transpose(1, 2).reshape(bsz, h, w, -1) output = output + lepe output = x + self.out_proj(output) output = output + self.layer_norm(output) return output def reset_parameters(self): nn.init.xavier_normal_(self.q_proj.weight, gain=2**-2.5) nn.init.xavier_normal_(self.k_proj.weight, gain=2**-2.5) nn.init.xavier_normal_(self.v_proj.weight, gain=2**-2.5) nn.init.xavier_normal_(self.out_proj.weight) nn.init.constant_(self.out_proj.bias, 0.0) class DepthCoordConv(nn.Module): """CoordConv with one additional normalized depth channel.""" def __init__( self, in_channels, out_channels, kernel_size=3, padding=1, stride=1, normalize_depth=True, eps=1e-6, ): super().__init__() self.normalize_depth = normalize_depth self.eps = eps self._depth = None self.conv1 = conv_layer( in_channels + 3, out_channels, kernel_size, padding, stride, ) def set_depth(self, depth): self._depth = depth def _resize_depth(self, depth, x): b, _, h, w = x.shape if depth is None: return x.new_zeros((b, 1, h, w)) if depth.dim() == 3: depth = depth.unsqueeze(1) depth = depth.to(device=x.device, dtype=x.dtype) if depth.shape[-2:] != (h, w): depth = F.interpolate( depth, size=(h, w), mode="bilinear", align_corners=False, ) depth = torch.nan_to_num(depth, nan=0.0, posinf=1.0, neginf=0.0) if self.normalize_depth: valid = depth > 0 valid_min = depth.masked_fill(~valid, float("inf")).amin( dim=(2, 3), keepdim=True, ) valid_max = depth.masked_fill(~valid, float("-inf")).amax( dim=(2, 3), keepdim=True, ) has_valid = valid.any(dim=(2, 3), keepdim=True) depth_min = torch.where(has_valid, valid_min, torch.zeros_like(valid_min)) depth_max = torch.where(has_valid, valid_max, torch.ones_like(valid_max)) depth = (depth - depth_min) / (depth_max - depth_min + self.eps) depth = torch.where(valid, depth, torch.zeros_like(depth)) return depth.clamp(0.0, 1.0) def add_coord(self, x, depth=None): b, _, h, w = x.shape x_range = torch.linspace(-1, 1, w, device=x.device, dtype=x.dtype) y_range = torch.linspace(-1, 1, h, device=x.device, dtype=x.dtype) y_grid, x_grid = torch.meshgrid(y_range, x_range) x_grid = x_grid.expand(b, 1, h, w) y_grid = y_grid.expand(b, 1, h, w) depth = self._resize_depth(depth, x) return torch.cat([x, x_grid, y_grid, depth], dim=1) def forward(self, x, depth=None): if depth is None: depth = self._depth x = self.add_coord(x, depth) return self.conv1(x)