| 253 | |
| 254 | |
| 255 | class PDNorm(PointModule): |
| 256 | def __init__( |
| 257 | self, |
| 258 | num_features, |
| 259 | norm_layer, |
| 260 | context_channels=256, |
| 261 | conditions=("ScanNet", "S3DIS", "Structured3D"), |
| 262 | decouple=True, |
| 263 | adaptive=False, |
| 264 | ): |
| 265 | super().__init__() |
| 266 | self.conditions = conditions |
| 267 | self.decouple = decouple |
| 268 | self.adaptive = adaptive |
| 269 | if self.decouple: |
| 270 | self.norm = nn.ModuleList([norm_layer(num_features) for _ in conditions]) |
| 271 | else: |
| 272 | self.norm = norm_layer |
| 273 | if self.adaptive: |
| 274 | self.modulation = nn.Sequential( |
| 275 | nn.SiLU(), nn.Linear(context_channels, 2 * num_features, bias=True) |
| 276 | ) |
| 277 | |
| 278 | def forward(self, point): |
| 279 | assert {"feat", "condition"}.issubset(point.keys()) |
| 280 | if isinstance(point.condition, str): |
| 281 | condition = point.condition |
| 282 | else: |
| 283 | condition = point.condition[0] |
| 284 | if self.decouple: |
| 285 | assert condition in self.conditions |
| 286 | norm = self.norm[self.conditions.index(condition)] |
| 287 | else: |
| 288 | norm = self.norm |
| 289 | point.feat = norm(point.feat) |
| 290 | if self.adaptive: |
| 291 | assert "context" in point.keys() |
| 292 | shift, scale = self.modulation(point.context).chunk(2, dim=1) |
| 293 | point.feat = point.feat * (1.0 + scale) + shift |
| 294 | return point |
| 295 | |
| 296 | |
| 297 | class RPE(torch.nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected