Extend functionality of torch.nn.functional.fold to apply to sequences (1D tensors + batch & channels) not just images (2D tensors + batch & channels)

[tom-hehir]

🚀 The feature, motivation and pitch

torch.nn.functional.unfold can be applied to sequences (1D tensors + batch & channels) and it would be helpful to be able to reverse this operation (up to addition in overlapping regions which can be addressed via a division -- see documentation for torch.nn.functional.fold) using torch.nn.functional.fold, however torch.nn.functional.fold only works on images (2D tensors + batch & channels) so this is not possible at present.
I suggest that torch.nn.functional.fold is generalised so that it can reverse torch.nn.functional.unfold in general (or at least for sequences), not just for images.

Alternatives

I have found a way to do this using Tensor.index_add (sketched out below):

a = torch.tensor(range(1,18))
print(a)

b = a.unfold(0, 5, 2)

out = torch.zeros_like(a)
indices = b.flatten()-1
out = out.index_add(0, indices, b.flatten())
x = out

a = torch.ones((17,))
b = a.unfold(0, 5, 2)
out = torch.zeros_like(a)
out = out.index_add(0, indices, b.flatten())
y = out

z = x / y
print(z)

It would be nice to abstract these details away and just call torch.nn.functiona.fold, and I expect that this kind of approach is far from optimal.

Additional context

No response

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki

[Microno95]

In case it's of utility, based on your implementation using Tensor.index_add, I've implemented an N-Dimensional Fold/Unfold module. It's more general than simply folding/unfolding as it allows you to modify the Unfold'ed blocks and then Fold them using new parameters (e.g. in the case where you sum the whole block and are left with blocks of size 1).

class NDUnfold(torch.nn.Module):
    def __init__(self, ndim, kernel_size, padding, stride):
        super().__init__()
        if isinstance(kernel_size, int):
            kernel_size = tuple([kernel_size]*ndim)
        if isinstance(padding, int):
            padding = tuple([padding]*ndim)
        if isinstance(stride, int):
            stride = tuple([stride]*ndim)
        padding = [min(max(k//2, 0), p) for k, p in zip(kernel_size, padding)]
        self.kernel_size = kernel_size
        self.padding = padding
        self.stride = stride
        self.repr_str = ", ".join(map(str, [ndim, tuple(kernel_size), tuple(padding), tuple(stride)]))

    def forward(self, in_tensor):
        ndim = len(self.kernel_size)
        x = torch.nn.functional.pad(in_tensor, [pad for padding in zip(self.padding, self.padding) for pad in padding], mode='constant', value=0)
        # print(in_tensor.shape, "->", x.shape)
        if x.dim() == ndim + 1:
            x = x.permute(*range(1, ndim+1), 0)
            offset = 0
        else:
            x = x.permute(0, *range(2, ndim+2), 1)
            offset = 1
        for idx in range(ndim):
            x = x.unfold(idx+offset, self.kernel_size[idx], self.stride[idx])
        x = x.flatten(start_dim=-ndim-1, end_dim=-1)
        x = x.flatten(start_dim=offset, end_dim=offset + ndim - 1)
        if x.dim() == 3:
            x = x.permute(0, 2, 1)
        else:
            x = x.permute(1, 0)
        # print("folded_shape:", x.shape, "from", in_tensor.shape)
        return x

    def __repr__(self):
        return f"{self.__class__.__name__}({self.repr_str})"
        
class NDFold(torch.nn.Module):
    def __init__(self, ndim, output_channels, output_kernel_size, output_padding, output_stride, input_kernel_size, input_size, input_padding, input_stride):
        super().__init__()
        if isinstance(output_kernel_size, int):
            output_kernel_size = tuple([output_kernel_size]*ndim)
        if isinstance(input_kernel_size, int):
            input_kernel_size = tuple([input_kernel_size]*ndim)
        if isinstance(input_size, int):
            input_size = tuple([input_size]*ndim)
        if isinstance(output_padding, int):
            output_padding = tuple([output_padding]*ndim)
        if isinstance(output_stride, int):
            output_stride = tuple([output_stride]*ndim)
        if isinstance(input_padding, int):
            input_padding = tuple([input_padding]*ndim)
        if isinstance(input_stride, int):
            input_stride = tuple([input_stride]*ndim)
        
        input_padding = [min(max(k//2, 0), p) for k, p in zip(input_kernel_size, input_padding)]
        
        self.ndim = ndim
        self.final_size = [int(round((s + 2*p - (k - 1) - 1)/stp+0.5)) for s, p, k, stp in zip(input_size, input_padding, input_kernel_size, input_stride)]
        self.output_channels = output_channels
        self.repr_str = ", ".join(map(str, [ndim, output_channels, tuple(output_kernel_size), tuple(output_padding), tuple(output_stride), tuple(input_kernel_size), tuple(input_size), tuple(input_padding), tuple(input_stride)]))
        
        # print("final_dims:", self.final_size, "| num_spatial:", int(np.prod(self.final_size)), "| num_channels:", self.output_channels)
        self.fold_weights = torch.nn.Parameter(NDUnfold(ndim, output_kernel_size, output_padding, output_stride)(torch.ones(self.output_channels, *self.final_size)), requires_grad=False)
        self.fold_indices = torch.nn.Parameter(NDUnfold(ndim, output_kernel_size, output_padding, output_stride)(torch.arange(0, int(self.output_channels*np.prod(self.final_size))).unflatten(0, (self.output_channels, *self.final_size))), requires_grad=False)
        # print("index shapes:", self.fold_indices.shape, "|", self.fold_weights.shape)

    def forward(self, in_tensor: torch.Tensor):
        ndim = self.ndim
        fold_weights = self.fold_weights
        fold_indices = self.fold_indices
        if in_tensor.dim() == 3:
            x = in_tensor.permute(0, 2, 1)
            out = torch.zeros(x.size(0), self.output_channels, *self.final_size, dtype=x.dtype, device=x.device)
        else:
            x = in_tensor.permute(1, 0)
            out = torch.zeros(self.output_channels, *self.final_size, dtype=x.dtype, device=x.device)
        weight_local = torch.zeros(self.output_channels, *self.final_size, dtype=x.dtype, device=x.device).flatten()
        weight_local = weight_local.index_add(0, self.fold_indices.flatten(), self.fold_weights.flatten())
        if in_tensor.dim() == 3:
            out_local = out.flatten(start_dim=1, end_dim=-1)
            out_local = out_local.index_add(1, self.fold_indices.flatten(), x.mT.flatten(start_dim=1, end_dim=-1))
            out = torch.unflatten(out_local / weight_local, 1, out.shape[1:])
        else:
            out_local = out.flatten()
            out_local = out_local.index_add(0, self.fold_indices.flatten(), x.mT.flatten())
            out = torch.unflatten(out_local / weight_local, 0, out.shape)
        return out

    def __repr__(self):
        return f"{self.__class__.__name__}({self.repr_str})"

I've tested it by checking that the Unfold operation is identical to torch.nn.Unfold when given an image input and that the two operations are inverses of each other:

test_x = torch.linspace(-1.0, 1.0, 8)
test_tensor = test_x[None, :, None, None].expand(3, -1, 4, 7)
print(NDUnfold(2, 3, 1, 1)(test_tensor).shape)
print(torch.allclose(torch.nn.Unfold(3, 1, padding=1, stride=1)(test_tensor), NDUnfold(2, 3, 1, 1)(test_tensor)))
test_tensor = test_x[None, :, None].expand(3, -1, 4)
print(torch.allclose(test_tensor, NDFold(1, test_x.shape[0], 3, 1, 1, 3, test_tensor.shape[2:], 1, 1)(NDUnfold(1, 3, 1, 1)(test_tensor))))
test_tensor = test_x[None, :, None, None].expand(3, -1, 4, 8)
print(torch.allclose(test_tensor, NDFold(test_tensor.dim()-2, test_x.shape[0], 3, 1, 1, 3, test_tensor.shape[2:], 1, 1)(NDUnfold(test_tensor.dim()-2, 3, 1, 1)(test_tensor))))
test_tensor = test_x[None, :, None, None, None].expand(3, -1, 4, 8, 16)
print(torch.allclose(test_tensor, NDFold(test_tensor.dim()-2, test_x.shape[0], 3, 1, 1, 3, test_tensor.shape[2:], 1, 1)(NDUnfold(test_tensor.dim()-2, 3, 1, 1)(test_tensor))))
test_tensor = test_x[None, :, None, None, None].expand(3, -1, 4, 8, 16)
print(torch.allclose(test_tensor, NDFold(test_tensor.dim()-2, test_x.shape[0], 1, 1, 1, 1, test_tensor.shape[2:], 1, 1)(NDUnfold(test_tensor.dim()-2, 1, 1, 1)(test_tensor))))

I realise it's a module-based implementation rather than a function-based one and that this could probably be improved further, but hopefully it's useful to anyone who finds their way to this issue. Disclaimer: I can't guarantee it works with whatever you throw at it