| # MIT License | |
| # Copyright (c) Microsoft Corporation. | |
| # Copyright (c) 2025 VAST-AI-Research and contributors. | |
| # Permission is hereby granted, free of charge, to any person obtaining a copy | |
| # of this software and associated documentation files (the "Software"), to deal | |
| # in the Software without restriction, including without limitation the rights | |
| # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell | |
| # copies of the Software, and to permit persons to whom the Software is | |
| # furnished to do so, subject to the following conditions: | |
| # The above copyright notice and this permission notice shall be included in all | |
| # copies or substantial portions of the Software. | |
| # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR | |
| # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, | |
| # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE | |
| # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER | |
| # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, | |
| # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE | |
| # SOFTWARE | |
| from typing import * | |
| import torch | |
| import torch.nn as nn | |
| from . import SparseTensor | |
| import torch.nn.functional as F | |
| import spconv.pytorch as spconv | |
| from typing import Optional | |
| from query_point import PE_NeRF | |
| from ...modules.sparse.transformer.blocks import SparseTransformerCrossBlock | |
| __all__ = [ | |
| 'SparseDownsample', | |
| 'SparseUpsample', | |
| 'SparseSubdivide', | |
| 'SparseSubdivide_attn' | |
| ] | |
| class SparseDownsample(nn.Module): | |
| """ | |
| Downsample a sparse tensor by a factor of `factor`. | |
| Implemented as average pooling. | |
| """ | |
| def __init__(self, factor: Union[int, Tuple[int, ...], List[int]]): | |
| super(SparseDownsample, self).__init__() | |
| self.factor = tuple(factor) if isinstance(factor, (list, tuple)) else factor | |
| def forward(self, input: SparseTensor) -> SparseTensor: | |
| DIM = input.coords.shape[-1] - 1 | |
| factor = self.factor if isinstance(self.factor, tuple) else (self.factor,) * DIM | |
| assert DIM == len(factor), 'Input coordinates must have the same dimension as the downsample factor.' | |
| coord = list(input.coords.unbind(dim=-1)) | |
| for i, f in enumerate(factor): | |
| coord[i+1] = coord[i+1] // f | |
| MAX = [coord[i+1].max().item() + 1 for i in range(DIM)] | |
| OFFSET = torch.cumprod(torch.tensor(MAX[::-1]), 0).tolist()[::-1] + [1] | |
| code = sum([c * o for c, o in zip(coord, OFFSET)]) | |
| code, idx = code.unique(return_inverse=True) | |
| new_feats = torch.scatter_reduce( | |
| torch.zeros(code.shape[0], input.feats.shape[1], device=input.feats.device, dtype=input.feats.dtype), | |
| dim=0, | |
| index=idx.unsqueeze(1).expand(-1, input.feats.shape[1]), | |
| src=input.feats, | |
| # reduce='mean' | |
| reduce='amax', | |
| ) | |
| new_coords = torch.stack( | |
| [code // OFFSET[0]] + | |
| [(code // OFFSET[i+1]) % MAX[i] for i in range(DIM)], | |
| dim=-1 | |
| ) | |
| out = SparseTensor(new_feats, new_coords, input.shape,) | |
| out._scale = tuple([s // f for s, f in zip(input._scale, factor)]) | |
| out._spatial_cache = input._spatial_cache | |
| out.register_spatial_cache(f'upsample_{factor}_coords', input.coords) | |
| out.register_spatial_cache(f'upsample_{factor}_layout', input.layout) | |
| out.register_spatial_cache(f'upsample_{factor}_idx', idx) | |
| return out | |
| # class SparseDownsample(nn.Module): | |
| # """ | |
| # Downsample a sparse tensor by a factor of `factor`. | |
| # Implemented as average pooling. | |
| # """ | |
| # def __init__(self, factor: Union[int, Tuple[int, ...], List[int]]): | |
| # super(SparseDownsample, self).__init__() | |
| # self.factor = tuple(factor) if isinstance(factor, (list, tuple)) else factor | |
| # def forward(self, input: SparseTensor) -> SparseTensor: | |
| # DIM = input.coords.shape[-1] - 1 | |
| # factor = self.factor if isinstance(self.factor, tuple) else (self.factor,) * DIM | |
| # assert DIM == len(factor), 'Input coordinates must have the same dimension as the downsample factor.' | |
| # coord = list(input.coords.unbind(dim=-1)) | |
| # for i, f in enumerate(factor): | |
| # coord[i+1] = coord[i+1] // f | |
| # MAX = [coord[i+1].max().item() + 1 for i in range(DIM)] | |
| # OFFSET = torch.cumprod(torch.tensor(MAX[::-1]), 0).tolist()[::-1] + [1] | |
| # code = sum([c * o for c, o in zip(coord, OFFSET)]) | |
| # code, idx = code.unique(return_inverse=True) | |
| # new_feats = torch.scatter_reduce( | |
| # torch.zeros(code.shape[0], input.feats.shape[1], device=input.feats.device, dtype=input.feats.dtype), | |
| # dim=0, | |
| # index=idx.unsqueeze(1).expand(-1, input.feats.shape[1]), | |
| # src=input.feats, | |
| # reduce='mean' | |
| # ) | |
| # new_coords = torch.stack( | |
| # [code // OFFSET[0]] + | |
| # [(code // OFFSET[i+1]) % MAX[i] for i in range(DIM)], | |
| # dim=-1 | |
| # ) | |
| # out = SparseTensor(new_feats, new_coords, input.shape,) | |
| # out._scale = tuple([s // f for s, f in zip(input._scale, factor)]) | |
| # out._spatial_cache = input._spatial_cache | |
| # if out.get_spatial_cache(f'upsample_{factor}_coords') is not None: | |
| # out.register_spatial_cache(f'upsample_{factor}_coords', [*out.get_spatial_cache(f'upsample_{factor}_coords'), input.coords]) | |
| # out.register_spatial_cache(f'upsample_{factor}_layout', [*out.get_spatial_cache(f'upsample_{factor}_layout'), input.layout]) | |
| # out.register_spatial_cache(f'upsample_{factor}_idx', [*out.get_spatial_cache(f'upsample_{factor}_idx'), idx]) | |
| # else: | |
| # out.register_spatial_cache(f'upsample_{factor}_coords', [input.coords]) | |
| # out.register_spatial_cache(f'upsample_{factor}_layout', [input.layout]) | |
| # out.register_spatial_cache(f'upsample_{factor}_idx', [idx]) | |
| # return out | |
| class SparseUpsample(nn.Module): | |
| """ | |
| Upsample a sparse tensor by a factor of `factor`. | |
| Implemented as nearest neighbor interpolation. | |
| """ | |
| def __init__(self, factor: Union[int, Tuple[int, int, int], List[int]]): | |
| super(SparseUpsample, self).__init__() | |
| self.factor = tuple(factor) if isinstance(factor, (list, tuple)) else factor | |
| def forward(self, input: SparseTensor) -> SparseTensor: | |
| DIM = input.coords.shape[-1] - 1 | |
| factor = self.factor if isinstance(self.factor, tuple) else (self.factor,) * DIM | |
| assert DIM == len(factor), 'Input coordinates must have the same dimension as the upsample factor.' | |
| new_coords = input.get_spatial_cache(f'upsample_{factor}_coords') | |
| new_layout = input.get_spatial_cache(f'upsample_{factor}_layout') | |
| idx = input.get_spatial_cache(f'upsample_{factor}_idx') | |
| # print(len(new_coords)) | |
| new_coords = new_coords.pop(-1) | |
| new_layout = new_layout.pop(-1) | |
| idx = idx.pop(-1) | |
| if any([x is None for x in [new_coords, new_layout, idx]]): | |
| raise ValueError('Upsample cache not found. SparseUpsample must be paired with SparseDownsample.') | |
| new_feats = input.feats[idx] | |
| out = SparseTensor(new_feats, new_coords, input.shape, new_layout) | |
| out._scale = tuple([s * f for s, f in zip(input._scale, factor)]) | |
| out._spatial_cache = input._spatial_cache | |
| return out | |
| class SparseSubdivide(nn.Module): | |
| """ | |
| Upsample a sparse tensor by a factor of `factor`. | |
| Implemented as nearest neighbor interpolation. | |
| """ | |
| def __init__(self): | |
| super(SparseSubdivide, self).__init__() | |
| def forward(self, input: SparseTensor) -> SparseTensor: | |
| DIM = input.coords.shape[-1] - 1 | |
| # upsample scale=2^DIM | |
| n_cube = torch.ones([2] * DIM, device=input.device, dtype=torch.int) | |
| n_coords = torch.nonzero(n_cube) | |
| n_coords = torch.cat([torch.zeros_like(n_coords[:, :1]), n_coords], dim=-1) | |
| factor = n_coords.shape[0] | |
| assert factor == 2 ** DIM | |
| # print(n_coords.shape) | |
| new_coords = input.coords.clone() | |
| new_coords[:, 1:] *= 2 | |
| new_coords = new_coords.unsqueeze(1) + n_coords.unsqueeze(0).to(new_coords.dtype) | |
| new_feats = input.feats.unsqueeze(1).expand(input.feats.shape[0], factor, *input.feats.shape[1:]) | |
| out = SparseTensor(new_feats.flatten(0, 1), new_coords.flatten(0, 1), input.shape) | |
| out._scale = input._scale * 2 | |
| out._spatial_cache = input._spatial_cache | |
| return out | |
| #################### new ######################## | |
| # 730 add ca, | |
| class SparseSubdivide_attn(nn.Module): | |
| """ | |
| Attention-based upsampling: Compute child voxel features with multi-head attention | |
| Enhanced with residual connections, layer normalization, and position encoding | |
| Improvements to overcome training plateau: | |
| 1. Position encoding using relative offsets instead of indices | |
| 2. Feature enhancement before attention | |
| 3. Output normalization and projection | |
| 4. Careful residual connections | |
| """ | |
| def __init__(self, in_channels: int, num_heads: int = 4,): | |
| super().__init__() | |
| assert in_channels % num_heads == 0, "in_channels must be divisible by num_heads" | |
| self.in_channels = in_channels | |
| self.head_dim = in_channels // num_heads | |
| self.num_heads = num_heads | |
| self.scale = self.head_dim ** -0.5 | |
| # Enhanced position encoding (continuous offsets instead of discrete indices) | |
| # self.pos_embed = nn.Sequential( | |
| # nn.Linear(3, 64), # Process actual spatial offsets | |
| # nn.LayerNorm(64), | |
| # nn.GELU(), | |
| # nn.Linear(64, in_channels) # Map to feature dimension | |
| # ) | |
| self.pos_embed = nn.Sequential( | |
| PE_NeRF(out_channels=in_channels, multires=10), # Process actual spatial offsets | |
| nn.LayerNorm(in_channels * 3), | |
| nn.GELU(), | |
| nn.Linear(in_channels * 3, in_channels) # Map to feature dimension | |
| ) | |
| # Feature enhancement before attention | |
| self.feat_enhance = nn.Sequential( | |
| nn.Linear(in_channels, in_channels * 2), | |
| nn.LayerNorm(in_channels * 2), | |
| nn.GELU(), | |
| nn.Linear(in_channels * 2, in_channels), | |
| nn.LayerNorm(in_channels) | |
| ) | |
| # Attention projections | |
| self.q_proj = nn.Linear(in_channels, in_channels) # Query from position | |
| self.k_proj = nn.Linear(in_channels, in_channels) # Key from content | |
| self.v_proj = nn.Linear(in_channels, in_channels) # Value from content | |
| # Output processing with residual | |
| self.output_norm = nn.LayerNorm(in_channels) | |
| self.output_proj = nn.Sequential( | |
| nn.Linear(in_channels, in_channels * 2), | |
| nn.GELU(), | |
| nn.Linear(in_channels * 2, in_channels) | |
| ) | |
| # Initialize for stable training | |
| self._initialize_weights() | |
| def _initialize_weights(self): | |
| nn.init.xavier_uniform_(self.q_proj.weight) | |
| nn.init.xavier_uniform_(self.k_proj.weight) | |
| nn.init.xavier_uniform_(self.v_proj.weight) | |
| nn.init.zeros_(self.q_proj.bias) | |
| nn.init.zeros_(self.k_proj.bias) | |
| nn.init.zeros_(self.v_proj.bias) | |
| nn.init.zeros_(self.output_proj[-1].weight) | |
| nn.init.zeros_(self.output_proj[-1].bias) | |
| nn.init.uniform_(self.output_proj[-1].weight, -1e-5, 1e-5) | |
| nn.init.constant_(self.output_proj[-1].bias, 0) | |
| def forward(self, input: SparseTensor) -> SparseTensor: | |
| DIM = input.coords.shape[-1] - 1 # Spatial dimensions (3 for 3D) | |
| device = input.device | |
| batch_coords = input.coords | |
| feats = input.feats | |
| # Generate child positions (identical to original) | |
| n_cube = torch.ones([2] * DIM, device=device, dtype=torch.int) | |
| n_coords = torch.nonzero(n_cube) # [8, 3] for 3D | |
| n_coords = torch.cat([torch.zeros_like(n_coords[:, :1]), n_coords], dim=-1) | |
| # Calculate actual spatial offsets (normalized to [-0.5, 0.5]) | |
| spatial_offsets = (n_coords[:, 1:].float() - 0.5) # Centered at origin | |
| pos_emb = self.pos_embed(spatial_offsets) # [8, C] | |
| # Enhance original features before attention | |
| enhanced_feats = self.feat_enhance(feats) | |
| # Compute new coordinates (same as original) | |
| new_coords = batch_coords.clone() | |
| new_coords[:, 1:] *= 2 | |
| expanded_coords = new_coords.unsqueeze(1) + n_coords.unsqueeze(0).to(new_coords.dtype) | |
| # Prepare attention inputs | |
| N = feats.shape[0] # Number of parent voxels | |
| num_children = n_coords.shape[0] # Always 8 for 3D | |
| # Project features to K, V | |
| K = self.k_proj(enhanced_feats).view(N, 1, self.num_heads, self.head_dim) | |
| V = self.v_proj(enhanced_feats).view(N, 1, self.num_heads, self.head_dim) | |
| # Project position embeddings to Q | |
| Q = self.q_proj(pos_emb).view(1, num_children, self.num_heads, self.head_dim) | |
| ######################################### | |
| # # Expand tensors for attention | |
| # K = K.expand(-1, num_children, -1, -1) # [N, 8, H, D] | |
| # V = V.expand(-1, num_children, -1, -1) # [N, 8, H, D] | |
| # Q = Q.expand(N, -1, -1, -1) # [N, 8, H, D] | |
| # attn_out = F.scaled_dot_product_attention( | |
| # Q.permute(0, 2, 1, 3).reshape(N * num_children, self.num_heads, self.head_dim), | |
| # K.permute(0, 2, 1, 3).reshape(N * num_children, self.num_heads, self.head_dim), | |
| # V.permute(0, 2, 1, 3).reshape(N * num_children, self.num_heads, self.head_dim), | |
| # dropout_p=0.0, | |
| # ) | |
| # # Reshape attention output | |
| # attn_out = attn_out.view(N, num_children, self.in_channels) | |
| K = K.expand(-1, num_children, -1, -1) # [N, 8, H, D] | |
| V = V.expand(-1, num_children, -1, -1) # [N, 8, H, D] | |
| Q = Q.expand(N, num_children, -1, -1) # [N, 8, H, D] | |
| # === 手动 scaled dot-product attention === | |
| Q_ = Q.permute(0, 2, 1, 3) # [N, H, 8, D] | |
| K_ = K.permute(0, 2, 1, 3) # [N, H, 8, D] | |
| V_ = V.permute(0, 2, 1, 3) # [N, H, 8, D] | |
| scale = self.head_dim ** -0.5 | |
| attn_logits = torch.matmul(Q_, K_.transpose(-2, -1)) * scale # [N, H, 8, 8] | |
| # 稳定 softmax：减去最大值 | |
| attn_logits = attn_logits - attn_logits.amax(dim=-1, keepdim=True) | |
| attn_weights = torch.softmax(attn_logits, dim=-1) | |
| attn_weights = torch.nan_to_num(attn_weights, nan=0.0, posinf=0.0, neginf=0.0) | |
| attn_output = torch.matmul(attn_weights, V_) # [N, H, 8, D] | |
| # 拼回原始形状 [N, 8, C] | |
| attn_out = attn_output.permute(0, 2, 1, 3).reshape(N, num_children, self.in_channels) | |
| # Position injection and output processing | |
| modulated = attn_out + pos_emb.unsqueeze(0) # Inject position info | |
| transformed = self.output_proj(self.output_norm(modulated)) | |
| # Residual connection: Combine with expanded parent features | |
| base_features = enhanced_feats.unsqueeze(1).expand(-1, num_children, -1) | |
| child_feats = base_features + transformed # Preserve original information | |
| # Create new sparse tensor | |
| out = SparseTensor( | |
| child_feats.reshape(N * num_children, -1), | |
| expanded_coords.flatten(0, 1), | |
| input.shape | |
| ) | |
| out._scale = input._scale * 2 | |
| out._spatial_cache = input._spatial_cache | |
| return out | |
| # ######################## relative linear ############################# | |
| # class SparseSubdivide_attn(nn.Module): | |
| # def __init__(self, in_channels: int, num_heads: int = 4, dropout: float = 0.05): | |
| # super().__init__() | |
| # assert in_channels % num_heads == 0, "in_channels must be divisible by num_heads" | |
| # self.in_channels = in_channels | |
| # self.head_dim = in_channels // num_heads | |
| # self.num_heads = num_heads | |
| # self.scale = self.head_dim ** -0.5 | |
| # self.pos_embed = nn.Sequential( | |
| # nn.Linear(3, in_channels), | |
| # nn.LayerNorm(in_channels), | |
| # nn.GELU(), | |
| # nn.Linear(in_channels, in_channels) | |
| # ) | |
| # self.feat_enhance = nn.Sequential( | |
| # nn.Linear(in_channels, in_channels * 2), | |
| # nn.LayerNorm(in_channels * 2), | |
| # nn.GELU(), | |
| # nn.Linear(in_channels * 2, in_channels), | |
| # nn.LayerNorm(in_channels) | |
| # ) | |
| # self.q_proj = nn.Linear(in_channels, in_channels) | |
| # self.k_proj = nn.Linear(in_channels, in_channels) | |
| # self.v_proj = nn.Linear(in_channels, in_channels) | |
| # self.output_norm = nn.LayerNorm(in_channels) | |
| # self.output_proj = nn.Sequential( | |
| # nn.Linear(in_channels, in_channels * 2), | |
| # nn.GELU(), | |
| # nn.Dropout(dropout), | |
| # nn.Linear(in_channels * 2, in_channels) | |
| # ) | |
| # # self._initialize_weights() | |
| # def _initialize_weights(self): | |
| # nn.init.xavier_uniform_(self.q_proj.weight) | |
| # nn.init.xavier_uniform_(self.k_proj.weight) | |
| # nn.init.xavier_uniform_(self.v_proj.weight) | |
| # nn.init.zeros_(self.q_proj.bias) | |
| # nn.init.zeros_(self.k_proj.bias) | |
| # nn.init.zeros_(self.v_proj.bias) | |
| # nn.init.zeros_(self.output_proj[-1].weight) | |
| # nn.init.constant_(self.output_proj[-1].bias, 0) | |
| # def forward(self, input: SparseTensor) -> SparseTensor: | |
| # DIM = input.coords.shape[-1] - 1 | |
| # device = input.device | |
| # coords = input.coords | |
| # feats = input.feats | |
| # N = feats.shape[0] | |
| # n_cube = torch.ones([2] * DIM, device=device, dtype=torch.int) | |
| # n_coords = torch.nonzero(n_cube) # [8, 3] | |
| # n_coords = torch.cat([torch.zeros_like(n_coords[:, :1]), n_coords], dim=-1) # [8, 4] | |
| # spatial_offsets = n_coords[:, 1:].float() # [8, 3], 用于位置编码 | |
| # pos_emb = self.pos_embed(spatial_offsets.to(device=device, dtype=feats.dtype)) # [8, C] | |
| # Q = self.q_proj(pos_emb).view(1, 8, self.num_heads, self.head_dim).expand(N, -1, -1, -1) # [N, 8, H, D] | |
| # enhanced_feats = self.feat_enhance(feats) # [N, C] | |
| # K = self.k_proj(enhanced_feats).view(N, 1, self.num_heads, self.head_dim) # [N, 1, H, D] | |
| # V = self.v_proj(enhanced_feats).view(N, 1, self.num_heads, self.head_dim) # [N, 1, H, D] | |
| # Q_ = Q.permute(0, 2, 1, 3) # [N, H, 8, D] | |
| # K_ = K.permute(0, 2, 1, 3) # [N, H, 1, D] | |
| # V_ = V.permute(0, 2, 1, 3) # [N, H, 1, D] | |
| # attn_logits = torch.matmul(Q_, K_.transpose(-2, -1)) * self.scale # [N, H, 8, 1] | |
| # attn_weights = torch.softmax(attn_logits, dim=2) # over children | |
| # attn_weights = torch.nan_to_num(attn_weights, nan=0.0, posinf=0.0, neginf=0.0) | |
| # attn_output = torch.matmul(attn_weights, V_) # [N, H, 8, D] | |
| # attn_out = attn_output.permute(0, 2, 1, 3).reshape(N, 8, self.in_channels) # [N, 8, C] | |
| # modulated = attn_out + pos_emb.unsqueeze(0) # [N, 8, C] | |
| # transformed = self.output_proj(self.output_norm(modulated)) # [N, 8, C] | |
| # base_features = enhanced_feats.unsqueeze(1).expand(-1, 8, -1) | |
| # child_feats = base_features + transformed # [N, 8, C] | |
| # new_coords = coords.clone() | |
| # new_coords[:, 1:] *= 2 | |
| # expanded_coords = new_coords.unsqueeze(1) + n_coords.to(dtype=coords.dtype).unsqueeze(0) # [N, 8, 4] | |
| # return SparseTensor( | |
| # child_feats.reshape(N * 8, self.in_channels), | |
| # expanded_coords.reshape(N * 8, 4), | |
| # input.shape | |
| # ) | |
| # ######################## relative embedding ############################# | |
| # class SparseSubdivide_attn(nn.Module): | |
| # def __init__(self, in_channels: int, num_heads: int = 4, dropout: float = 0.05): | |
| # super().__init__() | |
| # assert in_channels % num_heads == 0, "in_channels must be divisible by num_heads" | |
| # self.in_channels = in_channels | |
| # self.head_dim = in_channels // num_heads | |
| # self.num_heads = num_heads | |
| # self.scale = self.head_dim ** -0.5 | |
| # self.pos_index_embed = nn.Embedding(8, in_channels) | |
| # self.feat_enhance = nn.Sequential( | |
| # nn.Linear(in_channels, in_channels * 2), | |
| # nn.LayerNorm(in_channels * 2), | |
| # nn.GELU(), | |
| # nn.Linear(in_channels * 2, in_channels), | |
| # nn.LayerNorm(in_channels) | |
| # ) | |
| # self.q_proj = nn.Linear(in_channels, in_channels) | |
| # self.k_proj = nn.Linear(in_channels, in_channels) | |
| # self.v_proj = nn.Linear(in_channels, in_channels) | |
| # self.output_norm = nn.LayerNorm(in_channels) | |
| # self.output_proj = nn.Sequential( | |
| # nn.Linear(in_channels, in_channels * 2), | |
| # nn.GELU(), | |
| # nn.Dropout(dropout), | |
| # nn.Linear(in_channels * 2, in_channels) | |
| # ) | |
| # def forward(self, input: SparseTensor) -> SparseTensor: | |
| # DIM = input.coords.shape[-1] - 1 | |
| # device = input.device | |
| # coords = input.coords | |
| # feats = input.feats | |
| # N = feats.shape[0] | |
| # n_cube = torch.ones([2] * DIM, device=device, dtype=torch.int) | |
| # n_coords = torch.nonzero(n_cube) | |
| # n_coords = torch.cat([torch.zeros_like(n_coords[:, :1]), n_coords], dim=-1) # [8, 4] | |
| # pos_indices = torch.arange(8, device=device) | |
| # pos_emb = self.pos_index_embed(pos_indices) # [8, C] | |
| # Q = self.q_proj(pos_emb).view(1, 8, self.num_heads, self.head_dim).expand(N, -1, -1, -1) # [N, 8, H, D] | |
| # enhanced_feats = self.feat_enhance(feats) # [N, C] | |
| # K = self.k_proj(enhanced_feats).view(N, 1, self.num_heads, self.head_dim) # [N, 1, H, D] | |
| # V = self.v_proj(enhanced_feats).view(N, 1, self.num_heads, self.head_dim) # [N, 1, H, D] | |
| # # === Cross Attention === | |
| # Q_ = Q.permute(0, 2, 1, 3) # [N, H, 8, D] | |
| # K_ = K.permute(0, 2, 1, 3) # [N, H, 1, D] | |
| # V_ = V.permute(0, 2, 1, 3) # [N, H, 1, D] | |
| # attn_logits = torch.matmul(Q_, K_.transpose(-2, -1)) * self.scale # [N, H, 8, 1] | |
| # attn_weights = torch.softmax(attn_logits, dim=2) | |
| # attn_weights = torch.nan_to_num(attn_weights, nan=0.0, posinf=0.0, neginf=0.0) | |
| # attn_output = torch.matmul(attn_weights, V_) # [N, H, 8, D] | |
| # attn_out = attn_output.permute(0, 2, 1, 3).reshape(N, 8, self.in_channels) # [N, 8, C] | |
| # modulated = attn_out + pos_emb.unsqueeze(0) # [N, 8, C] | |
| # transformed = self.output_proj(self.output_norm(modulated)) # [N, 8, C] | |
| # base_features = enhanced_feats.unsqueeze(1).expand(-1, 8, -1) | |
| # child_feats = base_features + transformed # [N, 8, C] | |
| # new_coords = coords.clone() | |
| # new_coords[:, 1:] *= 2 | |
| # expanded_coords = new_coords.unsqueeze(1) + n_coords.to(dtype=coords.dtype).unsqueeze(0) # [N, 8, 4] | |
| # return SparseTensor( | |
| # child_feats.reshape(N * 8, self.in_channels), | |
| # expanded_coords.reshape(N * 8, 4), | |
| # input.shape | |
| # ) | |
| # ############################## Position-Specific Filters ##################################### | |
| # class SparseSubdivide_attn(nn.Module): | |
| # def __init__(self, in_channels: int, num_heads: int = 4, dropout: float = 0.05): | |
| # super().__init__() | |
| # self.in_channels = in_channels | |
| # self.num_heads = num_heads | |
| # self.head_dim = in_channels // num_heads | |
| # self.scale = self.head_dim ** -0.5 | |
| # # Position-aware modulation components | |
| # self.pos_index_embed = nn.Embedding(8, in_channels) | |
| # self.pos_filters = nn.Embedding(8, in_channels * in_channels) | |
| # # Feature enhancement | |
| # self.feat_enhance = nn.Sequential( | |
| # nn.Linear(in_channels, in_channels * 2), | |
| # nn.LayerNorm(in_channels * 2), | |
| # nn.GELU(), | |
| # nn.Linear(in_channels * 2, in_channels), | |
| # nn.LayerNorm(in_channels) | |
| # ) | |
| # # Output transformation | |
| # self.output_norm = nn.LayerNorm(in_channels) | |
| # self.output_proj = nn.Sequential( | |
| # nn.Linear(in_channels, in_channels * 2), | |
| # nn.GELU(), | |
| # nn.Dropout(dropout), | |
| # nn.Linear(in_channels * 2, in_channels) | |
| # ) | |
| # def forward(self, input: SparseTensor) -> SparseTensor: | |
| # DIM = input.coords.shape[-1] - 1 | |
| # device = input.device | |
| # coords = input.coords | |
| # feats = input.feats | |
| # N = feats.shape[0] | |
| # # Generate subdivision coordinates (8 children per voxel) | |
| # n_cube = torch.ones([2] * DIM, device=device, dtype=torch.int) | |
| # n_coords = torch.nonzero(n_cube) | |
| # n_coords = torch.cat([torch.zeros_like(n_coords[:, :1]), n_coords], dim=-1) # [8, 4] | |
| # pos_indices = torch.arange(8, device=device) | |
| # # Position-aware feature modulation | |
| # pos_emb = self.pos_index_embed(pos_indices) # [8, C] | |
| # pos_filters = self.pos_filters(pos_indices) # [8, C*C] | |
| # pos_filters = pos_filters.view(8, self.in_channels, self.in_channels) # [8, C, C] | |
| # # Enhance parent features | |
| # enhanced_feats = self.feat_enhance(feats) # [N, C] | |
| # # Apply position-specific transformation | |
| # modulated_feats = torch.einsum('pci,nc->npc', pos_filters, enhanced_feats) # [N, 8, C] | |
| # modulated_feats = modulated_feats + pos_emb.unsqueeze(0) # Add positional encoding | |
| # # Final transformation | |
| # transformed = self.output_proj(self.output_norm(modulated_feats)) # [N, 8, C] | |
| # child_feats = enhanced_feats.unsqueeze(1) + transformed # Residual connection | |
| # # Compute new coordinates | |
| # new_coords = coords.clone() | |
| # new_coords[:, 1:] *= 2 | |
| # expanded_coords = new_coords.unsqueeze(1) + n_coords.to(dtype=coords.dtype).unsqueeze(0) # [N, 8, 4] | |
| # return SparseTensor( | |
| # child_feats.reshape(N * 8, self.in_channels), | |
| # expanded_coords.reshape(N * 8, 4), | |
| # input.shape | |
| # ) | |
| # # ################################# nn.Parameter ####################################### | |
| # # class SparseSubdivideCrossAttn(nn.Module): | |
| # # def __init__(self, in_channels: int, num_heads: int = 4, mlp_ratio: int = 4): | |
| # # super().__init__() | |
| # # self.in_channels = in_channels | |
| # # self.num_heads = num_heads | |
| # # self.head_dim = in_channels // num_heads | |
| # # self.scale = self.head_dim ** -0.5 | |
| # # self.mlp_ratio = mlp_ratio | |
| # # self.pos_embed = nn.Parameter(torch.randn(8, in_channels)) | |
| # # self.q_proj = nn.Linear(in_channels, in_channels) | |
| # # self.kv_proj = nn.Linear(in_channels, in_channels * 2) | |
| # # self.proj = nn.Linear(in_channels, in_channels) | |
| # # self.norm1 = nn.LayerNorm(in_channels) | |
| # # self.norm2 = nn.LayerNorm(in_channels) | |
| # # self.mlp = nn.Sequential( | |
| # # nn.Linear(in_channels, in_channels * mlp_ratio), | |
| # # nn.GELU(), | |
| # # nn.Linear(in_channels * mlp_ratio, in_channels) | |
| # # ) | |
| # # def forward(self, input: SparseTensor) -> SparseTensor: | |
| # # DIM = input.coords.shape[-1] - 1 | |
| # # device = input.device | |
| # # coords = input.coords | |
| # # feats = input.feats | |
| # # N = feats.shape[0] | |
| # # n_cube = torch.ones([2] * DIM, device=device, dtype=torch.int) | |
| # # n_coords = torch.nonzero(n_cube) | |
| # # n_coords = torch.cat([torch.zeros_like(n_coords[:, :1]), n_coords], dim=-1) | |
| # # q = self.q_proj(self.pos_embed) | |
| # # q = q.reshape(8, self.num_heads, self.head_dim).permute(1, 0, 2) # [num_heads, 8, head_dim] | |
| # # kv = self.kv_proj(feats) | |
| # # kv = kv.reshape(N, 2, self.num_heads, self.head_dim).permute(2, 0, 1, 3) # [num_heads, N, 2, head_dim] | |
| # # k, v = kv[:, :, 0, :], kv[:, :, 1, :] # [num_heads, N, head_dim] for both | |
| # # attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale # [num_heads, 8, N] | |
| # # attn = torch.softmax(attn, dim=-1) | |
| # # out = torch.matmul(attn, v) # [num_heads, 8, head_dim] | |
| # # out = out.permute(1, 0, 2).reshape(8, N, self.in_channels) # [8, N, in_channels] | |
| # # out = out.permute(1, 0, 2) # [N, 8, in_channels] | |
| # # x = self.proj(out) + feats.unsqueeze(1) | |
| # # x = self.norm1(x) | |
| # # x = x + self.mlp(self.norm2(x)) | |
| # # new_coords = coords.clone() | |
| # # new_coords[:, 1:] *= 2 | |
| # # expanded_coords = new_coords.unsqueeze(1) + n_coords.to(dtype=coords.dtype).unsqueeze(0) | |
| # # return SparseTensor( | |
| # # x.reshape(N * 8, self.in_channels), | |
| # # expanded_coords.reshape(N * 8, 4), | |
| # # input.shape | |
| # # ) | |
| # # ################################ Modulation #################################### | |
| # # class SparseSubdivideModulation(nn.Module): | |
| # # def __init__(self, in_channels: int): | |
| # # super().__init__() | |
| # # self.in_channels = in_channels | |
| # # self.position_emb = nn.Embedding(8, in_channels) | |
| # # self.modulation_vectors = nn.Embedding(8, in_channels) | |
| # # self.feature_transformer = nn.Sequential( | |
| # # nn.Linear(in_channels, in_channels * 2), | |
| # # nn.LayerNorm(in_channels * 2), | |
| # # nn.GELU(), | |
| # # nn.Linear(in_channels * 2, in_channels) | |
| # # ) | |
| # # self.output_mlp = nn.Sequential( | |
| # # nn.Linear(in_channels, in_channels * 2), | |
| # # nn.GELU(), | |
| # # nn.Linear(in_channels * 2, in_channels) | |
| # # ) | |
| # # def forward(self, input: SparseTensor) -> SparseTensor: | |
| # # DIM = input.coords.shape[-1] - 1 | |
| # # device = input.device | |
| # # coords = input.coords | |
| # # feats = input.feats | |
| # # N = feats.shape[0] | |
| # # n_cube = torch.ones([2] * DIM, device=device, dtype=torch.int) | |
| # # n_coords = torch.nonzero(n_cube) | |
| # # n_coords = torch.cat([torch.zeros_like(n_coords[:, :1]), n_coords], dim=-1) | |
| # # pos_ids = torch.arange(8, device=device) | |
| # # trans_feats = self.feature_transformer(feats) # [N, C] | |
| # # mod_vectors = self.modulation_vectors(pos_ids) # [8, C] | |
| # # pos_emb = self.position_emb(pos_ids) # [8, C] | |
| # # modulated_feats = torch.einsum('nc,pc->npc', trans_feats, mod_vectors) + pos_emb | |
| # # output_feats = self.output_mlp(modulated_feats) # [N, 8, C] | |
| # # new_coords = coords.clone() | |
| # # new_coords[:, 1:] *= 2 | |
| # # expanded_coords = new_coords.unsqueeze(1) + n_coords.to(dtype=coords.dtype).unsqueeze(0) | |
| # # return SparseTensor( | |
| # # output_feats.reshape(N * 8, self.in_channels), | |
| # # expanded_coords.reshape(N * 8, 4), | |
| # # input.shape | |
| # # ) | |
| # ############################## 16 * 3 embedding ############################## | |
| # # 730 | |
| class SparseSubdivide_attn(nn.Module): | |
| def __init__(self, in_channels: int, num_heads: int = 4, resolution: int=128): | |
| super().__init__() | |
| self.in_channels = in_channels | |
| self.num_heads = num_heads | |
| self.embed_dim = in_channels | |
| self.relative_coords = resolution // 64 | |
| self.coord_embed_x = nn.Embedding(self.relative_coords, self.embed_dim) | |
| self.coord_embed_y = nn.Embedding(self.relative_coords, self.embed_dim) | |
| self.coord_embed_z = nn.Embedding(self.relative_coords, self.embed_dim) | |
| self.embed_proj = nn.Linear(in_channels * 3, in_channels) | |
| def forward(self, input): | |
| DIM = input.coords.shape[-1] - 1 | |
| device = input.device | |
| feats = input.feats | |
| n_cube = torch.ones([2]*DIM, device=device, dtype=torch.int) | |
| n_coords = torch.nonzero(n_cube) | |
| n_coords = torch.cat([torch.zeros_like(n_coords[:, :1]), n_coords], dim=-1) | |
| new_coords = input.coords.clone() | |
| new_coords[:, 1:] *= 2 | |
| new_coords = new_coords.unsqueeze(1) + n_coords.unsqueeze(0).to(new_coords.dtype) # [N, 8, 4] | |
| abs_coords = new_coords[:, :, 1:] # [N, 8, 3] | |
| mod_coords = abs_coords % self.relative_coords | |
| x_embed = self.coord_embed_x(mod_coords[..., 0].long()) | |
| y_embed = self.coord_embed_y(mod_coords[..., 1].long()) | |
| z_embed = self.coord_embed_z(mod_coords[..., 2].long()) | |
| pos_embed = torch.cat([x_embed, y_embed, z_embed], dim=-1) # [N, 8, 3C] | |
| pos_embed = self.embed_proj(pos_embed) # [N, 8, C] | |
| feats = feats.unsqueeze(1).expand(-1, 8, -1) # [N, 8, C] | |
| new_feats = feats + pos_embed # [N, 8, C] | |
| out = SparseTensor( | |
| new_feats.flatten(0, 1), | |
| new_coords.flatten(0, 1), | |
| input.shape | |
| ) | |
| out._scale = input._scale * 2 | |
| out._spatial_cache = input._spatial_cache | |
| return out | |
| class SparseSubdivide_attn(nn.Module): | |
| """ | |
| Upsample with sparse cross-attention between parent features and position embeddings | |
| """ | |
| def __init__(self, in_channels: int, num_heads: int = 4): | |
| super().__init__() | |
| self.in_channels = in_channels | |
| self.num_heads = num_heads | |
| self.pos_embed = nn.Embedding(8, in_channels) # [8, C] | |
| def forward(self, input: SparseTensor) -> SparseTensor: | |
| DIM = input.coords.shape[-1] - 1 | |
| device = input.device | |
| # upsample scale=2^DIM | |
| n_cube = torch.ones([2] * DIM, device=input.device, dtype=torch.int) | |
| n_coords = torch.nonzero(n_cube) | |
| n_coords = torch.cat([torch.zeros_like(n_coords[:, :1]), n_coords], dim=-1) | |
| factor = n_coords.shape[0] | |
| assert factor == 2 ** DIM | |
| # print(n_coords.shape) | |
| new_coords = input.coords.clone() | |
| new_coords[:, 1:] *= 2 | |
| new_coords = new_coords.unsqueeze(1) + n_coords.unsqueeze(0).to(new_coords.dtype) | |
| base_feats = input.feats.unsqueeze(1).expand(-1, factor, -1) # [N,8,C] | |
| child_ids = torch.arange(8, device=device) # [8] | |
| pos_feats = self.pos_embed(child_ids).unsqueeze(0) # [1,8,C] | |
| query_feats = base_feats + pos_feats # [N,8,C] | |
| final_feats = query_feats.flatten(0, 1) | |
| out = SparseTensor( | |
| feats=final_feats, | |
| coords=new_coords.flatten(0, 1), | |
| shape=input.shape | |
| ) | |
| out._scale = input._scale * 2 | |
| return out |