How can I fold a Tensor that I unfolded with PyTorch that has overlap?

Viewed 589

I have a Tensor of size: torch.Size([1, 63840]) which I then unrolled:

inp_unfolded = inp_seq.unfold(1, 160, 80)

that gives me a shape of: torch.Size([1, 797, 160])

How can I re-fold that to get a Tensor of torch.Size([1, 63840])?

2 Answers

For that specific configuration, since 63840 is divisible by 160 and the step size is a multiple of the slice size, you can simply select every second element along that dimension and then flatten the resulting tensor:

inp_unfolded[:, ::2, :].flatten(1, 2)

More generally, for t.unfold(i, n, s), if t.shape[i] % n == 0 and n % s == 0 holds, then you can restore the original tensor via:

index = [slice(None) for __ in t.shape]
index[i] = slice(None, None, n // s)
original = t.unfold(i, n, s)[tuple(index)].flatten(i, i+1)

Of course you can also use slice notation, if the dimension i is known beforehand. For example i == 1 as in your example:

original = t.unfold(1, n, s)[:, ::n//s, ...].flatten(1, 2)

Well, actually the conditions, given t.unfold(i, n, s) are:

  • n >= s (otherwise step is skipping some original data and we cannot restore it)
  • n + s <= t.shape[i]

Then we can do it via:

def roll(x, n, s, axis=1):
    return torch.cat((p[0], p[1:][:, n-s:].flatten()), axis)

explanation:

p[0] is the starting chunk that is always unique at start

p[1:][:, n-s:] - then, we take rest of rolls and n-s depict how many elements will overlap between rolls so we want to ignore them and take only those from n-s

ilustration:

x.unfold(0, 5, 2)
tensor([[ 1.,  2.,  3.,  4.,  5.],
        [ 3.,  4.,  5.,  6.,  7.], # 3, 4, 5 are repeated
        [ 5.,  6.,  7.,  8.,  9.], # 5, 6, 7 are repeated...
        [ 7.,  8.,  9., 10., 11.],
        [ 9., 10., 11., 12., 13.],
        [11., 12., 13., 14., 15.],
        [13., 14., 15., 16., 17.]])

example:

>> x = torch.arange(1., 18)
>> p = x.unfold(0, 5, 2)
>> roll(p, 5, 2, 0)

tensor([ 1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10., 11., 12., 13., 14.,
        15., 16., 17.])

you can also try it with

x = torch.arange(1., 18).reshape(1, 17)

and axis 1

Related