- added size_splits to functional

[ptrblck]

Pull request addresses issue #3223

The split function splits tensors into equally sized chunks.
split_sizes let the user define a list with sizes for each chunk.

tf.split combines both functionalities in one function. Maybe this is also desired for Pytorch?

split_sizes seems to be a bit slower (6.991s vs 6.704s)

[colesbury]

Yeah, I think a single function would be nicer. Both tf.split and numpy.split have that sort of API.

[ptrblck]

Thanks for the feedback, I will merge it in torch.split then.

[ptrblck]

I merged both functions now.

Before that, I timed all functions on my machine (tensor=torch.randn(200, 10, 2, 2) into 40 chunks of size 5 in dim=0)

• merged function with split_size_or_sections = [5] * 40: ~8.1586s
• merged function with split_size_or_sections = 5: ~7.4982
• native with `split_size=5: ~7.2688s

What are your thoughts? Any suggestion on the code / naming / documentation?

[ezyang]

@pytorchbot test this please

[flennerhag]

@ptrblck great initiative, I've been using my own wrapper for a while. Would be nice to have it in the code base.

You could simplify the code quite a bit though by using plain python operations instead of invoking torch overheads, e.g. something like

def split(tensor, sizes, dim=0):
    if dim < 0:
        dim += tensor.dim()

    if isinstance(sizes, int):
        # original code ...
        return chunks

    if tensor.size(dim) != sum(sizes):
        raise ValueError("Sizes do not match tensor size in dim")

    nsizes = len(sizes)
    sizes = [0] + sizes
    return tuple(tensor.narrow(dim, sizes[i], sizes[i + 1])
                 for i in range(nsizes))

Should be slightly faster too.

[soumith]

@ptrblck as soon as you add unit tests for the list of splits case, i can merge this in.

[ptrblck]

@flennerhag Thanks for the suggestions! I tried to change some Pytorch code to plain python operations.
@soumith I also added some tests in test_split. If it's not sufficient, I can add some more test cases.

[soumith]

@pytorchbot test this please

[soumith]

thanks a lot @ptrblck !