(self, feats, num_patches=64, patch_ids=None)
| 586 | self.mlp_init = True |
| 587 | |
| 588 | def forward(self, feats, num_patches=64, patch_ids=None): |
| 589 | return_ids = [] |
| 590 | return_feats = [] |
| 591 | if self.use_mlp and not self.mlp_init: |
| 592 | self.create_mlp(feats) |
| 593 | for feat_id, feat in enumerate(feats): |
| 594 | B, H, W = feat.shape[0], feat.shape[2], feat.shape[3] |
| 595 | feat_reshape = feat.permute(0, 2, 3, 1).flatten(1, 2) |
| 596 | if num_patches > 0: |
| 597 | if patch_ids is not None: |
| 598 | patch_id = patch_ids[feat_id] |
| 599 | else: |
| 600 | patch_id = torch.randperm(feat_reshape.shape[1], device=feats[0].device) |
| 601 | patch_id = patch_id[:int(min(num_patches, patch_id.shape[0]))] # .to(patch_ids.device) |
| 602 | x_sample = feat_reshape[:, patch_id, :].flatten(0, 1) # reshape(-1, x.shape[1]) |
| 603 | else: |
| 604 | x_sample = feat_reshape |
| 605 | patch_id = [] |
| 606 | if self.use_mlp: |
| 607 | mlp = getattr(self, 'mlp_%d' % feat_id) |
| 608 | x_sample = mlp(x_sample) |
| 609 | return_ids.append(patch_id) |
| 610 | x_sample = self.l2norm(x_sample) |
| 611 | |
| 612 | if num_patches == 0: |
| 613 | x_sample = x_sample.permute(0, 2, 1).reshape([B, x_sample.shape[-1], H, W]) |
| 614 | return_feats.append(x_sample) |
| 615 | return return_feats, return_ids |
| 616 | |
| 617 | |
| 618 | class G_Resnet(nn.Module): |
nothing calls this directly
no test coverage detected