File size: 5,022 Bytes
97fd768 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 | """Standalone architecture extracted from the original lidarx implementation."""
import torch
from torch import nn
def coarse_projection_features(x: torch.Tensor) -> torch.Tensor:
"""Return `projection_coarse_32x8x8` features for `[B, 1, D, H, W]` tensors."""
if x.ndim != 5:
msg = f"expected input with shape [batch, channel, depth, height, width], got {x.shape}"
raise ValueError(msg)
if x.shape[1] != 1:
msg = f"expected one histogram channel, got {x.shape[1]}"
raise ValueError(msg)
x = x.squeeze(1)
depth = x.shape[1]
height = x.shape[2]
width = x.shape[3]
if depth % 32 != 0 or height % 8 != 0 or width % 8 != 0:
msg = (
"hist coarse projection requires depth divisible by 32 and "
f"height/width divisible by 8, got {(depth, height, width)}"
)
raise ValueError(msg)
global_features = torch.cat(
[
x.mean(dim=(1, 2, 3), keepdim=False).unsqueeze(1),
x.std(dim=(1, 2, 3), unbiased=False, keepdim=False).unsqueeze(1),
x.amin(dim=(1, 2, 3), keepdim=False).unsqueeze(1),
x.amax(dim=(1, 2, 3), keepdim=False).unsqueeze(1),
],
dim=1,
)
depth_features = torch.cat(
[
x.mean(dim=(2, 3)),
x.std(dim=(2, 3), unbiased=False),
],
dim=1,
)
row_features = x.mean(dim=(1, 3))
col_features = x.mean(dim=(1, 2))
coarse = (
x.reshape(
x.shape[0],
32,
depth // 32,
8,
height // 8,
8,
width // 8,
)
.mean(dim=(2, 4, 6))
.reshape(x.shape[0], -1)
)
return torch.cat(
[global_features, depth_features, row_features, col_features, coarse],
dim=1,
)
def _coarse_grid_features(
x: torch.Tensor,
*,
depth_bins: int,
row_bins: int,
col_bins: int,
) -> torch.Tensor:
depth = x.shape[1]
height = x.shape[2]
width = x.shape[3]
if depth % depth_bins != 0 or height % row_bins != 0 or width % col_bins != 0:
msg = (
"hist coarse grid requires dimensions divisible by requested bins, "
f"got volume {(depth, height, width)} and bins {(depth_bins, row_bins, col_bins)}"
)
raise ValueError(msg)
return (
x.reshape(
x.shape[0],
depth_bins,
depth // depth_bins,
row_bins,
height // row_bins,
col_bins,
width // col_bins,
)
.mean(dim=(2, 4, 6))
.reshape(x.shape[0], -1)
)
def multiscale_coarse_projection_features(x: torch.Tensor) -> torch.Tensor:
"""Return coarse histogram features with added low-resolution context views."""
base = coarse_projection_features(x)
volume = x.squeeze(1)
depth_context = _coarse_grid_features(volume, depth_bins=16, row_bins=8, col_bins=8)
spatial_context = _coarse_grid_features(volume, depth_bins=32, row_bins=4, col_bins=4)
return torch.cat([base, depth_context, spatial_context], dim=1)
class HistCoarseMultiScaleMLP(nn.Module):
"""Small supervised MLP over multi-scale coarse histogram projections.
This keeps the strong non-backbone inductive bias of the coarse-only model,
but gives the classifier both the best 32x8x8 block view and two coarser
context views before a lightweight gated residual head.
"""
def __init__(self, hidden_features: int = 128, dropout: float = 0.2) -> None:
super().__init__()
bottleneck_features = max(1, hidden_features // 2)
self.input_norm = nn.LazyBatchNorm1d()
self.input_projection = nn.Sequential(
nn.LazyLinear(hidden_features),
nn.SiLU(inplace=True),
nn.Dropout(dropout),
)
self.residual = nn.Sequential(
nn.Linear(hidden_features, hidden_features),
nn.SiLU(inplace=True),
nn.Dropout(dropout),
nn.Linear(hidden_features, hidden_features),
)
self.gate = nn.Sequential(
nn.Linear(hidden_features, hidden_features),
nn.Sigmoid(),
)
self.head = nn.Sequential(
nn.LayerNorm(hidden_features),
nn.SiLU(inplace=True),
nn.Dropout(dropout),
nn.Linear(hidden_features, bottleneck_features),
nn.SiLU(inplace=True),
nn.Dropout(dropout),
nn.Linear(bottleneck_features, 1),
)
def forward_features(self, x: torch.Tensor) -> torch.Tensor:
"""Return the learned 128-D representation immediately before the classifier."""
features = self.input_norm(multiscale_coarse_projection_features(x))
embedding = self.input_projection(features)
return embedding + self.gate(embedding) * self.residual(embedding)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.head(self.forward_features(x))
|