| import torch |
| import torch.nn as nn |
| import numpy as np |
|
|
|
|
| class AbstractPermuter(nn.Module): |
| def __init__(self, *args, **kwargs): |
| super().__init__() |
| def forward(self, x, reverse=False): |
| raise NotImplementedError |
|
|
|
|
| class Identity(AbstractPermuter): |
| def __init__(self): |
| super().__init__() |
|
|
| def forward(self, x, reverse=False): |
| return x |
|
|
|
|
| class Subsample(AbstractPermuter): |
| def __init__(self, H, W): |
| super().__init__() |
| C = 1 |
| indices = np.arange(H*W).reshape(C,H,W) |
| while min(H, W) > 1: |
| indices = indices.reshape(C,H//2,2,W//2,2) |
| indices = indices.transpose(0,2,4,1,3) |
| indices = indices.reshape(C*4,H//2, W//2) |
| H = H//2 |
| W = W//2 |
| C = C*4 |
| assert H == W == 1 |
| idx = torch.tensor(indices.ravel()) |
| self.register_buffer('forward_shuffle_idx', |
| nn.Parameter(idx, requires_grad=False)) |
| self.register_buffer('backward_shuffle_idx', |
| nn.Parameter(torch.argsort(idx), requires_grad=False)) |
|
|
| def forward(self, x, reverse=False): |
| if not reverse: |
| return x[:, self.forward_shuffle_idx] |
| else: |
| return x[:, self.backward_shuffle_idx] |
|
|
|
|
| def mortonify(i, j): |
| """(i,j) index to linear morton code""" |
| i = np.uint64(i) |
| j = np.uint64(j) |
|
|
| z = np.uint(0) |
|
|
| for pos in range(32): |
| z = (z | |
| ((j & (np.uint64(1) << np.uint64(pos))) << np.uint64(pos)) | |
| ((i & (np.uint64(1) << np.uint64(pos))) << np.uint64(pos+1)) |
| ) |
| return z |
|
|
|
|
| class ZCurve(AbstractPermuter): |
| def __init__(self, H, W): |
| super().__init__() |
| reverseidx = [np.int64(mortonify(i,j)) for i in range(H) for j in range(W)] |
| idx = np.argsort(reverseidx) |
| idx = torch.tensor(idx) |
| reverseidx = torch.tensor(reverseidx) |
| self.register_buffer('forward_shuffle_idx', |
| idx) |
| self.register_buffer('backward_shuffle_idx', |
| reverseidx) |
|
|
| def forward(self, x, reverse=False): |
| if not reverse: |
| return x[:, self.forward_shuffle_idx] |
| else: |
| return x[:, self.backward_shuffle_idx] |
|
|
|
|
| class SpiralOut(AbstractPermuter): |
| def __init__(self, H, W): |
| super().__init__() |
| assert H == W |
| size = W |
| indices = np.arange(size*size).reshape(size,size) |
|
|
| i0 = size//2 |
| j0 = size//2-1 |
|
|
| i = i0 |
| j = j0 |
|
|
| idx = [indices[i0, j0]] |
| step_mult = 0 |
| for c in range(1, size//2+1): |
| step_mult += 1 |
| |
| for k in range(step_mult): |
| i = i - 1 |
| j = j |
| idx.append(indices[i, j]) |
|
|
| |
| for k in range(step_mult): |
| i = i |
| j = j + 1 |
| idx.append(indices[i, j]) |
|
|
| step_mult += 1 |
| if c < size//2: |
| |
| for k in range(step_mult): |
| i = i + 1 |
| j = j |
| idx.append(indices[i, j]) |
|
|
| |
| for k in range(step_mult): |
| i = i |
| j = j - 1 |
| idx.append(indices[i, j]) |
| else: |
| |
| for k in range(step_mult-1): |
| i = i + 1 |
| idx.append(indices[i, j]) |
|
|
| assert len(idx) == size*size |
| idx = torch.tensor(idx) |
| self.register_buffer('forward_shuffle_idx', idx) |
| self.register_buffer('backward_shuffle_idx', torch.argsort(idx)) |
|
|
| def forward(self, x, reverse=False): |
| if not reverse: |
| return x[:, self.forward_shuffle_idx] |
| else: |
| return x[:, self.backward_shuffle_idx] |
|
|
|
|
| class SpiralIn(AbstractPermuter): |
| def __init__(self, H, W): |
| super().__init__() |
| assert H == W |
| size = W |
| indices = np.arange(size*size).reshape(size,size) |
|
|
| i0 = size//2 |
| j0 = size//2-1 |
|
|
| i = i0 |
| j = j0 |
|
|
| idx = [indices[i0, j0]] |
| step_mult = 0 |
| for c in range(1, size//2+1): |
| step_mult += 1 |
| |
| for k in range(step_mult): |
| i = i - 1 |
| j = j |
| idx.append(indices[i, j]) |
|
|
| |
| for k in range(step_mult): |
| i = i |
| j = j + 1 |
| idx.append(indices[i, j]) |
|
|
| step_mult += 1 |
| if c < size//2: |
| |
| for k in range(step_mult): |
| i = i + 1 |
| j = j |
| idx.append(indices[i, j]) |
|
|
| |
| for k in range(step_mult): |
| i = i |
| j = j - 1 |
| idx.append(indices[i, j]) |
| else: |
| |
| for k in range(step_mult-1): |
| i = i + 1 |
| idx.append(indices[i, j]) |
|
|
| assert len(idx) == size*size |
| idx = idx[::-1] |
| idx = torch.tensor(idx) |
| self.register_buffer('forward_shuffle_idx', idx) |
| self.register_buffer('backward_shuffle_idx', torch.argsort(idx)) |
|
|
| def forward(self, x, reverse=False): |
| if not reverse: |
| return x[:, self.forward_shuffle_idx] |
| else: |
| return x[:, self.backward_shuffle_idx] |
|
|
|
|
| class Random(nn.Module): |
| def __init__(self, H, W): |
| super().__init__() |
| indices = np.random.RandomState(1).permutation(H*W) |
| idx = torch.tensor(indices.ravel()) |
| self.register_buffer('forward_shuffle_idx', idx) |
| self.register_buffer('backward_shuffle_idx', torch.argsort(idx)) |
|
|
| def forward(self, x, reverse=False): |
| if not reverse: |
| return x[:, self.forward_shuffle_idx] |
| else: |
| return x[:, self.backward_shuffle_idx] |
|
|
|
|
| class AlternateParsing(AbstractPermuter): |
| def __init__(self, H, W): |
| super().__init__() |
| indices = np.arange(W*H).reshape(H,W) |
| for i in range(1, H, 2): |
| indices[i, :] = indices[i, ::-1] |
| idx = indices.flatten() |
| assert len(idx) == H*W |
| idx = torch.tensor(idx) |
| self.register_buffer('forward_shuffle_idx', idx) |
| self.register_buffer('backward_shuffle_idx', torch.argsort(idx)) |
|
|
| def forward(self, x, reverse=False): |
| if not reverse: |
| return x[:, self.forward_shuffle_idx] |
| else: |
| return x[:, self.backward_shuffle_idx] |
|
|
|
|
| if __name__ == "__main__": |
| p0 = AlternateParsing(16, 16) |
| print(p0.forward_shuffle_idx) |
| print(p0.backward_shuffle_idx) |
|
|
| x = torch.randint(0, 768, size=(11, 256)) |
| y = p0(x) |
| xre = p0(y, reverse=True) |
| assert torch.equal(x, xre) |
|
|
| p1 = SpiralOut(2, 2) |
| print(p1.forward_shuffle_idx) |
| print(p1.backward_shuffle_idx) |
|
|
