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

Yuerno: <p>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.</p>
<p>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).</p>
<p>For additional context, this is how I generate my 3D grid right now:</p>
<pre><code class="lang-auto">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))
</code></pre>

ptrblck: <p><a href="https://pytorch.org/docs/stable/tensors.html#torch.Tensor.unfold"><code>tensor.unfold</code></a> could be used here.<br>
Have a look at <a href="https://discuss.pytorch.org/t/creating-nonoverlapping-patches-from-3d-data-and-reshape-them-back-to-the-image/51210/6">this post</a> for an example.</p>

Yuerno: <p>Oh awesome, I think this might be exactly what I need. Thanks!</p>

Yuerno: <p>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.</p>

ptrblck: <p>You could initialize the input tensor with:</p>
<pre><code class="lang-python">x = torch.arange(batch_size*C*H*W)..view(batch_size, C, H, W).float()
</code></pre>
<p>so that you can visually inspect each patch and make sure it contains the right values.</p>
<p>I’m not sure I understand the last sentence. Would you like to fold them back to the original input afterwards?</p>

Yuerno: <p>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?</p>
<p>The context here is that I’m trying to match up a coordinates from a separate smaller grid with each block.</p>