Args: x: point features of shape [B, C, 3, N, ...] Returns: rotation invariant features of the same shape
(self, x)
| 270 | in_channels // 4, 2 if self.use_rmat else 3, dim=dim) |
| 271 | |
| 272 | def forward(self, x): |
| 273 | """ |
| 274 | Args: |
| 275 | x: point features of shape [B, C, 3, N, ...] |
| 276 | Returns: |
| 277 | rotation invariant features of the same shape |
| 278 | """ |
| 279 | z = self.vn1(x) |
| 280 | z = self.vn2(z) |
| 281 | z = self.vn_lin(z) # [B, 3, 3, N] or [B, 2, 3, N] |
| 282 | if self.use_rmat: |
| 283 | z = z.flatten(1, 2).transpose(1, 2).contiguous() # [B, N, 6] |
| 284 | z = rot6d_to_matrix(z) # [B, N, 3, 3] |
| 285 | z = z.permute(0, 2, 3, 1) # [B, 3, 3, N] |
| 286 | z = z.transpose(1, 2).contiguous() |
| 287 | |
| 288 | if self.dim == 4: |
| 289 | x_in = torch.einsum('bijm,bjkm->bikm', x, z) |
| 290 | elif self.dim == 3: |
| 291 | x_in = torch.einsum('bij,bjk->bik', x, z) |
| 292 | elif self.dim == 5: |
| 293 | x_in = torch.einsum('bijmn,bjkmn->bikmn', x, z) |
| 294 | else: |
| 295 | raise NotImplementedError(f'dim={self.dim} is not supported') |
| 296 | |
| 297 | return x_in |
| 298 | |
| 299 | |
| 300 | class VNLinearBNLeakyReLU(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected