Theme

Blog · Research infrastructure ·

A convolution that never voxelizes

Every algorithm I had just benchmarked turned a point cloud into a voxel grid first. I spent three and a half months trying to build one that does not, wired into the same ResNet and U-Net Chris Choy designed for MinkowskiEngine so the comparison would be fair. No result was ever recorded.

  • Interactive
  • pytorch
  • graph-neural-networks
  • point-clouds
  • geometric-deep-learning
  • research-infrastructure

Four algorithms, one benchmark was the payoff post of a six-post series about a harness I built to score other people’s 3D semantic segmentation fairly. MinkowskiEngine came out of that harness looking very good, and it, along with 3DMV, does the same thing before it does anything else: it turns the point cloud into a voxel grid. Every coordinate gets snapped onto a regular lattice, and the convolution runs over that lattice, sparse or not.

So I spent the six weeks after that benchmark trying to build the alternative: a convolution that gathers from a point’s actual neighbours, weighted by the direction to each one, and never puts the cloud on a grid at all. I wired it into the exact ResNet and U-Net topology MinkowskiEngine itself uses, so that if I ever got to compare the two, the comparison would be about the convolution and nothing else. This post is about that operator, why I built it the way I did, and the test I set up for it. It is not about a result, because no result was ever recorded. The project I built it in is not one I can link here (a private research repo, not one of the public ones this site usually cites), but every line quoted below is real code I wrote and ran.

The problem with a grid

The experiment script that first tried a graph convolution in place of MinkowskiEngine’s put the goal plainly in its own docstring:

EXPERIMENT_DESC = """--- AIM ---
To test the directed graph versions of resnet. This is done to see if it is possible to get rid of the ME convolution
which requires voxelisation. We also aim to see if directional information is able to assist a graph based approach.
--- PROCEDURE ---
1) Create a directional graph convolution layer
2) Replace the ME convolution layers and pooling layers with graph-based techniques
3) Test the accuracy, training speed and performance of the new model
--- HYPOTHESIS ---
The graph based approach has slightly fewer parameters but the way that they are used will enforce smoother kernals
and may even lead to an increase in accuracy. Performance may be slighly less due to the python implementation.
"""

A voxel grid is a fine idea when the thing you are scanning is roughly uniform, but a room is not: a wall is a dense flat sheet of points, the open middle of the floor is sparse, and a cluttered object sitting on that floor is denser again than either. Fit one cell size and it is either too coarse for the object or wastefully fine for the empty air around it. A convolution that reads from actual neighbours, not lattice offsets, does not have to choose.

The operator

The kernel weighting lives in calc_weighting, and it is small enough to read in one sitting:

def calc_weighting(pos, edges, flow='source_to_target', inc_center=True, max_dist=None, out=None):
    i, j = (0, 1) if flow == 'target_to_source' else (1, 0)
    sz = 7 if inc_center else 6
    out = out if out is not None else torch.empty(edges.shape[1], sz, dtype=torch.float, device=pos.device)
    out[:, :3] = pos[edges[j]] - pos[edges[i]]
    d_pos = torch.norm(out[:, :3], dim=1, keepdim=True) + 1e-10
    if inc_center:
        max_dist = d_pos.max() if max_dist is None else max_dist
        out[:, 6] = (max_dist - d_pos.view(-1)) / max_dist
    out[:, :3] /= d_pos
    out[:, 3:6] = -out[:, :3]
    torch.clamp_min(out, 0, out=out)
    # For version 3 have to transpose and add extra dim
    out = out.T[..., None]
    return out

For every edge, take the unit vector from the centre point to its neighbour. Columns 0-2 are that vector’s x, y, z components; columns 3-5 are the same vector negated. Clamp everything to zero and each neighbour ends up with a non-zero weight only in the columns that match the direction it actually lies in: a neighbour due east contributes only to the +x column, a neighbour to the northwest splits its weight between +y and -x, in proportion to how aligned it is with each axis. Column 6 is a separate, purely distance-based term: a linear falloff from 1 at the centre point itself to 0 at the edge of the search radius. Six axis bins plus one closeness bin, all built from nothing but the two points’ positions, no grid involved.

DirectedConv.message (directed_graph_convV2.py:71-78) then applies a learned in_channels × out_channels weight matrix per bin and sums:

def message(self, x_j, dir_weight):
    hold = torch.mul(dir_weight[0], x_j) @ self.weight[0]
    for i in range(1, self.weight_nr):
        hold += torch.mul(dir_weight[i], x_j) @ self.weight[i]
    return hold

and update (:80-88) adds the centre point’s own feature back in, through the same seventh weight matrix that the closeness bin already uses:

def update(self, inputs, x):
    if self.has_center:
        inputs = torch.addmm(inputs, x, self.weight[-1])
    if self.bias is not None:
        inputs += self.bias
    ...

That last line is something I only noticed rereading this for this post: weight[-1] is weight[6], the same matrix message already uses to gate a neighbour’s closeness contribution. Here it is reused, unweighted, against the centre point’s own feature. It is not a bug: it means the “self” term and the “how close was my nearest neighbour” term share one set of learned parameters, which is a smaller model than giving self its own eighth weight matrix, and I clearly meant to do it, since has_center is what turns the line on at all. It is exactly the kind of thing that is much easier to see two months after the fact than while writing it.

radius_graph (from torch_geometric, directed_graph_convV2.py:51-53) builds the edge list itself, with loop=False, from a max_dist I pass in as a constructor argument: each layer decides its own neighbourhood, exactly the role a MinkowskiEngine kernel size plays, but continuous instead of discretised. And the weighting is recomputed only when the graph actually changes:

if 'dir_weight' not in data or \
    data['dir_weight'] is None or \
    data['dir_weight'].shape[0] != data['edge_index'].shape[1] or \
    data['dir_weight'].shape[1] != self.weight_nr:

    data['dir_weight'] = calc_weighting(data['pos'], data['edge_index'], flow=self.flow,
                                        inc_center=self.has_center, max_dist=self.max_dist)\
        .to(data['x'].device)

computed once under torch.no_grad() and cached on the batch dict, keyed by nothing more than “does this shape still match the edge list.” A cheap guard, but it means a four-block ResNet layer that reuses the same radius does not recompute seven numbers per edge four times over.

Same topology, different convolution

Getting rid of the voxel grid is only half the point. The other half is that a fair comparison needs everything else held equal, so expbed/models/graphunet/dgresnet.py and dgUnet.py reproduce MinkowskiEngine’s own reference architecture almost line for line, down to the layer names, with ME.MinkowskiConvolution swapped for my DirectedConv. dgresnet.py::ResNetBase shares its class name with Choy’s own ResNetBase (resnet.py) and matches it structurally; dgUnet.py::Res16UNetBase does the same against Choy’s MinkUNetBase (minkunet.py) under a different class name but the exact same layer names and shape:

LayerMinkowskiEngine (resnet.py)Graph version (dgresnet.py)
stemself.conv1 = ME.MinkowskiConvolution(in_channels, self.inplanes, kernel_size=5, stride=2, dimension=D)layers['conv1'] = DirectedConv(in_channels, in_planes, max_dist=config.max_dist, groups=c_grp[0])
normME.MinkowskiBatchNormNormLayer(NormType.BATCH_NORM, ...)
stagesself.layer1 = self._make_layer(self.BLOCK, self.PLANES[0], self.LAYERS[0], stride=2)layers['layer1'], in_planes = self._make_layer(self.BLOCK, in_planes, self.PLANES[0], self.LAYERS[0], c_grp[1])
downsamplingstride on the sparse convolutiona separate ConvPool message-passing layer (aggr='add')

Same PLANES = (64, 128, 256, 512), same LAYERS, same block classes (BasicBlock, Bottleneck), same _make_layer shape. The one structural difference is that Minkowski folds stride into the convolution itself; a radius graph has no notion of stride, so downsampling has to be its own explicit pooling step. Everything else - channel counts, block counts, where a downsample projection is inserted - matches on purpose.

The U-Net version (dgUnet.py’s Res16UNetBase) goes further: conv0p1s1, conv1p1s2, block1 through block8, convtr4p16s2 through convtr7p2s2 are the exact layer names from MinkowskiEngine’s MinkUNetBase (minkunet.py), and Res16UNet14A2Small even reuses Choy’s PLANES/LAYERS convention for naming a specific configuration. Nothing about that naming is an accident: it is the only way to make “same architecture, different convolution” a claim you can check by diffing two files instead of taking on faith.

The convolution that can pretend to be a grid

By late May the single DirectedConv had grown into a small family, all sharing one calc_weighting-style scheme but disagreeing about how continuous the direction should be. The commit that introduced them says what the two most interesting ones were for:

e8063f9  2020-05-27  Added 4 conv types for testing. SnapFullConv Should be equivelent to MinkUNet in a sense

FullConv is the continuous operator above, seven weight matrices and clamped dot products. SnapFullConv (directed_graph_conv_20200527.py:115-127) keeps the same seven bins but snaps each neighbour’s direction to whichever one of the six axes it is closest to, with argmax instead of a soft split:

class SnapFullConv(FullConv):
    def message(self, x_j, dir_weight):
        hold = torch.empty(x_j.shape[0], self.out_channels, device=x_j.device, dtype=x_j.dtype)
        for i in range(self.num_dir_weights):
            msk = dir_weight == i
            hold[msk] = x_j[msk] @ self.weight[i]
        return hold

SnapDoubleConv goes one step further and snaps to the nearest of all 26 directions of a 3×3×3 neighbourhood (every non-zero offset in {-1,0,1}³), which is as close as a continuous-coordinate graph convolution can get to reproducing a dense voxel kernel’s exact footprint without actually voxelizing:

DIRECTION_ARRAY = torch.tensor([[x, y, z] for x in range(-1, 2)
                                for y in range(-1, 2)
                                for z in range(-1, 2)
                                if x != 0 or y != 0 or z != 0], dtype=torch.float32)

That is a genuinely useful middle ground for a head-to-head test: if the fully continuous operator and MinkowskiEngine disagree, you cannot easily tell whether the disagreement is about voxelizing at all or about something else entirely. Testing the snapped version first isolates the “different operator, same discretisation” question before touching the “no discretisation at all” one.

The test I set up, and never finished

Two experiment scripts a day apart set the actual comparison up. e20200601 establishes what MinkowskiEngine gets on a smaller Res16UNet14A2, the real ME.MinkowskiConvolution version. e20200602_00_eq_mink_graph.py, the next day, is the graph side of the same comparison, Res16UNet14A2Small with conv_type = "SnapDoubleConv" - so the actual equivalence test ran the 26-direction snapped operator, not the fully continuous one, for exactly the isolation reason above. Its stated hypothesis:

EXPERIMENT_NAME = "Test equivilent graph architecture"
EXPERIMENT_DESC = """--- AIM ---
To test if the graph version of Res16UNet14A2 achieves similar results to the original.
--- PROCEDURE ---
--- HYPOTHESIS ---
The architectres are mathematically equivelent and should perform the same if there are no bugs in the model.
"""

(Typos as written; I have not corrected “architectres” or “equivelent”.) That sentence is the whole thesis of this post: same topology, a convolution that should, in principle, produce the same feature maps a sparse voxel convolution would, at the same lattice resolution, without ever needing to voxelize the input to get there. “If there are no bugs in the model” is doing a lot of work in that sentence, and I never found out whether there were.

Try it: one output point’s neighbourhood, two ways

The best way I can show what “never voxelizes” actually buys you is to look at one point’s neighbourhood both ways at once. The room below is synthetic - a floor, two walls meeting at a corner, and one cluttered object - built with genuinely non-uniform density on purpose, because that is exactly where the argument lives. Drag the dot to move the query point; the left panel gathers its real neighbours and weights each one by direction, exactly as calc_weighting does; the right panel snaps the same point onto a lattice of the same scale and shows the fixed 3×3 (or, with “3D directions” on, 3×3×3) block of cells around it, the way a sparse voxel convolution would.

InteractiveDirectional kernel explorer

With JavaScript enabled this becomes two linked panels over a synthetic room-corner point cloud. The left one gathers a draggable query point’s actual neighbours within a radius and weights each by the direction to it, exactly as the directed graph convolution’s calc_weighting does. The right one snaps the same point onto a lattice of the same scale and shows the fixed block of cells a sparse voxel convolution would read from instead, occupied or not. A bar chart below shows the resulting per-direction weighted sum, and a toggle switches between a simplified 2D view (four direction bins) and the real 3D one (six, plus a query-height control).

Watch what happens when the query point sits near the edge of the cluttered object: the graph panel’s neighbour count and its weighted bars change smoothly as you drag, because every point keeps its own exact position and distance right up until the final mean. The voxel panel’s occupied-cell count jumps in steps, and the readout below it will often report a handful of points sharing a single cell - their exact positions already thrown away before any weight was ever applied. That collapse-before-weighting, not just “a different set of neighbours,” is the actual mechanism behind “never voxelizes.”

A sidebar: keeping a lab notebook as git tags

One more small piece worth two paragraphs rather than its own post. Every experiment script in this project calls get_next_tag_name before it does anything else (expbed/syncfn/__init__.py:24), which regex-matches existing git tags shaped like exp_{TIMESTAMP}_{EXP_NR}_{EXP_NAME} and computes the next run number, and tag_git_exp (:56) refuses to actually create a tag until git status reports a clean tree. It is a lab notebook implemented entirely as git history: every run gets a unique, sortable, greppable tag, and you cannot tag a run whose code you have not committed, which rules out the “which version of the model actually produced this number” question by construction.

That is a different answer to the same problem the config file is the experiment solved a few months earlier, where the output directory’s YAML dump was the experiment record, reconstructable without git at all. Tags need the repository to make sense of; a YAML dump does not care where it is read from. Neither is more “correct” - one keeps the record next to the code that made it, the other keeps it next to the numbers it produced - but it is a genuinely different choice, made a few months apart, on the same underlying problem.

What actually shipped

The directed graph convolution and its ResNet/U-Net wrapping are real, tested-by-being-run code, not a sketch: real experiment scripts imported and trained them, the OOM-recovery training loop and the ScanNet dataloading stack around them worked well enough to keep running for two and a half months. What did not ship is a number. I built the operator, I built the fair comparison, and I ran out of thesis before I got to read the result. A portfolio post can show the idea and the test; it should not pretend to show the answer it never got.