Best way to split a 3D grid of points into blocks?

Hi all. If I have a batch of uniform 3D grids of coordinate locations, with shape (for example), [1, 32, 32, 32, 3], what is the best way for me to split this up into multiple even chunks, so I could end up with something such as [1, 4096, 2, 2, 2, 3]? In other words, I’m splitting up that one big 32 x 32 x 32 cube where each point is an x, y, z coordinate location into 4096 smaller 2 x 2 x 2 cubes? Does a simple view operation make sense here, or would it throw off the coordinate values? I was looking into operations like torch.chunk, but they require a specific dimension to split along, which I’m not sure applies here.

My use case for this is that I have a smaller [1, 16, 16, 16, 3] cube, so I’m trying to match up points from this smaller shape into the corresponding cubes in the upsampled [1, 32, 32, 32, 3] shape (since a single coordinate point in the 16^3 shape corresponds to 8 points in the 32^3 shape).

For additional context, this is how I generate my 3D grid right now:

pxs = torch.linspace(-1, 1, 32)
pys = torch.linspace(-1, 1, 32)
pzs = torch.linspace(-1, 1, 32)

pxs = pxs.view(-1, 1, 1).expand(*shape).contiguous().view(size)
pys = pys.view(1, -1, 1).expand(*shape).contiguous().view(size)
pzs = pzs.view(1, 1, -1).expand(*shape).contiguous().view(size)
points = torch.stack([pxs, pys, pzs], dim=1)

grid_3d = torch.reshape(points, (32, 32, 32, 3))

tensor.unfold could be used here.
Have a look at this post for an example.

1 Like

Oh awesome, I think this might be exactly what I need. Thanks!

1 Like

I hope you don’t mind me following up on this question with a related one! Tensor unfolding + your example is definitely what I’m looking for, but I was wondering if you had any ideas on ways to visualize/verify the blocks so I can double check that the blocks include the correct sections that I want (especially to double check if I’m using fold in the proper order of dimensions), as well as how they join together to make the whole block.

You could initialize the input tensor with:

x = torch.arange(batch_size*C*H*W)..view(batch_size, C, H, W).float()

so that you can visually inspect each patch and make sure it contains the right values.

I’m not sure I understand the last sentence. Would you like to fold them back to the original input afterwards?

Thanks for your reply and help so far, I really appreciate it! To clarify my last sentence, is there a way to check in what order you combine each of the smaller blocks to generate the larger block? For example, when I split up my 32x32x32 grid into 8 smaller 16x16x16 grids, the returned tensor from the fold operations will just have essentially a list of the blocks, but I don’t know if those blocks were taken by going horizontally across the larger grid, or by going vertically. I could match these up by manually inspecting the values in each of the smaller blocks, and trying to check where they appear in the larger block to see, but that’s a bit tedious. Hopefully I explained that better?

The context here is that I’m trying to match up a coordinates from a separate smaller grid with each block.