| 296 | |
| 297 | |
| 298 | class jigsaw_patch: |
| 299 | def __init__(self, shift=5, group=2): |
| 300 | self.shift = shift |
| 301 | self.group = group |
| 302 | |
| 303 | def __call__(self, features): |
| 304 | batchsize = features.size(0) |
| 305 | dim = features.size(1) |
| 306 | features = features.view(batchsize, dim, -1) |
| 307 | |
| 308 | # Shift Operation |
| 309 | feature_random = torch.cat([features[:, :, self.shift:], features[:, :, :self.shift]], dim=2) |
| 310 | x = feature_random |
| 311 | |
| 312 | # Patch Shuffle Operation |
| 313 | try: |
| 314 | x = x.view(batchsize, dim, self.group, -1) |
| 315 | except: |
| 316 | x = torch.cat([x, x[:, -2:-1, :]], dim=1) |
| 317 | x = x.view(batchsize, self.group, -1, dim) |
| 318 | |
| 319 | x = torch.transpose(x, 2, 3).contiguous() |
| 320 | |
| 321 | x = x.view(batchsize, dim, -1) |
| 322 | x = x.view(batchsize, dim, 14, 14) |
| 323 | |
| 324 | return x |
| 325 | |