x : B M 3 mask: B M ----------------- feature_global : B C
(self, x, mask=None)
| 269 | ) |
| 270 | |
| 271 | def forward(self, x, mask=None): |
| 272 | """ |
| 273 | x : B M 3 |
| 274 | mask: B M |
| 275 | ----------------- |
| 276 | feature_global : B C |
| 277 | """ |
| 278 | |
| 279 | bs, n, _ = x.shape |
| 280 | device = x.device |
| 281 | |
| 282 | x_valid = self.first_mlp(x[mask]) # B n 256 |
| 283 | x_features = torch.zeros(bs, n, 256, device=device) |
| 284 | x_features[mask] = x_valid |
| 285 | |
| 286 | pooled_feature = x_features.max(dim=1)[0] |
| 287 | x_features = torch.cat( |
| 288 | [x_features, pooled_feature.unsqueeze(1).repeat(1, n, 1)], dim=-1 |
| 289 | ) |
| 290 | |
| 291 | x_features_valid = self.second_mlp(x_features[mask]) |
| 292 | res = torch.zeros(bs, n, self.encoder_channel, device=device) |
| 293 | res[mask] = x_features_valid |
| 294 | |
| 295 | res = res.max(dim=1)[0] |
| 296 | return res |
nothing calls this directly
no outgoing calls
no test coverage detected