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:
| Layer | MinkowskiEngine (resnet.py) | Graph version (dgresnet.py) |
|---|---|---|
| stem | self.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]) |
| norm | ME.MinkowskiBatchNorm | NormLayer(NormType.BATCH_NORM, ...) |
| stages | self.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]) |
| downsampling | stride on the sparse convolution | a 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.
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.