| 564 | |
| 565 | |
| 566 | class PatchSampleF(nn.Module): |
| 567 | def __init__(self, use_mlp=False, init_type='normal', init_gain=0.02, nc=256, gpu_ids=[]): |
| 568 | # potential issues: currently, we use the same patch_ids for multiple images in the batch |
| 569 | super(PatchSampleF, self).__init__() |
| 570 | self.l2norm = Normalize(2) |
| 571 | self.use_mlp = use_mlp |
| 572 | self.nc = nc # hard-coded |
| 573 | self.mlp_init = False |
| 574 | self.init_type = init_type |
| 575 | self.init_gain = init_gain |
| 576 | self.gpu_ids = gpu_ids |
| 577 | |
| 578 | def create_mlp(self, feats): |
| 579 | for mlp_id, feat in enumerate(feats): |
| 580 | input_nc = feat.shape[1] |
| 581 | mlp = nn.Sequential(*[nn.Linear(input_nc, self.nc), nn.ReLU(), nn.Linear(self.nc, self.nc)]) |
| 582 | if len(self.gpu_ids) > 0: |
| 583 | mlp.cuda() |
| 584 | setattr(self, 'mlp_%d' % mlp_id, mlp) |
| 585 | init_net(self, self.init_type, self.init_gain, self.gpu_ids) |
| 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): |