Jigswa + linear + l2norm
| 13 | |
| 14 | |
| 15 | class JigsawHead(nn.Module): |
| 16 | """Jigswa + linear + l2norm""" |
| 17 | def __init__(self, dim_in, dim_out, k=9, head='linear'): |
| 18 | super(JigsawHead, self).__init__() |
| 19 | |
| 20 | if head == 'linear': |
| 21 | self.fc1 = nn.Linear(dim_in, dim_out) |
| 22 | elif head == 'mlp': |
| 23 | self.fc1 = nn.Sequential( |
| 24 | nn.Linear(dim_in, dim_in), |
| 25 | nn.ReLU(inplace=True), |
| 26 | nn.Linear(dim_in, dim_out), |
| 27 | ) |
| 28 | else: |
| 29 | raise NotImplementedError('JigSaw head not supported: {}'.format(head)) |
| 30 | self.fc2 = nn.Linear(dim_out * k, dim_out) |
| 31 | self.l2norm = Normalize(2) |
| 32 | self.k = k |
| 33 | |
| 34 | def forward(self, x): |
| 35 | bsz = x.shape[0] |
| 36 | x = self.fc1(x) |
| 37 | # ==== shuffle ==== |
| 38 | # this step can be moved to data processing step |
| 39 | shuffle_ids = self.get_shuffle_ids(bsz) |
| 40 | x = x[shuffle_ids] |
| 41 | # ==== shuffle ==== |
| 42 | n_img = int(bsz / self.k) |
| 43 | x = x.view(n_img, -1) |
| 44 | x = self.fc2(x) |
| 45 | x = self.l2norm(x) |
| 46 | return x |
| 47 | |
| 48 | def get_shuffle_ids(self, bsz): |
| 49 | n_img = int(bsz / self.k) |
| 50 | rnd_ids = [torch.randperm(self.k) for i in range(n_img)] |
| 51 | rnd_ids = torch.cat(rnd_ids, dim=0) |
| 52 | base_ids = torch.arange(bsz) |
| 53 | base_ids = torch.div(base_ids, self.k).long() |
| 54 | base_ids = base_ids * self.k |
| 55 | shuffle_ids = rnd_ids + base_ids |
| 56 | return shuffle_ids |