LLoCa-ParticleNet
We start with the standard ParticleNet implementation, available at https://github.com/hqucms/weaver-core/blob/main/weaver/nn/model/ParticleNet.py.
We implement LLoCa with the change_local_frame function that is applied in each message-passing step to transform the features from the sender frame to the receiver frame. We now also require the representation structure of the features, which has to be propagated through the architecture.
import torch
import torch.nn as nn
+from ..framesnet.frames import ChangeOfFrames, IndexSelectFrames
+from ..reps.tensorreps import TensorReps
+from ..reps.tensorreps_transform import TensorRepsTransform
+def change_local_frame(x_j_framej, idx, frames, trafo):
+ """Transform features x_j from frame 'j' ('x_j_framej') to frame 'i' ('x_j_framei').
+ Parameters
+ ----------
+ x_j_framej : torch.Tensor
+ Input features in local frame 'j' of shape (batch_size, num_dims, num_points, k).
+ idx : torch.Tensor
+ Indices of the nearest neighbors in the batch of shape (batch_size*num_points*k).
+ frames : Frames
+ Local frames of reference for the particles, shape (num_points, 4, 4).
+ trafo : TensorRepsTransform
+ Transformation function to apply to the features.
+ Returns
+ -------
+ torch.Tensor
+ """
+ # we use batch_size*num_points with repeats of k for idx_i, e.g. for 2 points with 3 batch and k=2,
+ # idx_i becomes (0,1,2,3,4,5) -> (0,0,1,1,2,2,3,3,4,4,5,5).
+ idx_i = torch.arange(
+ x_j_framej.shape[2] * x_j_framej.shape[0], device=x_j_framej.device
+ ).repeat_interleave(x_j_framej.shape[-1]) # identity (batch, num_points*k)
+ idx_j = idx # indices from knn (batch, num_points*k)
+ frames_i = IndexSelectFrames(frames, idx_i)
+ frames_j = IndexSelectFrames(frames, idx_j)
+ trafo_j_to_i = ChangeOfFrames(frames_j, frames_i) # convention: (frames_start, frames_end)
+ # reshape and apply trafo
+ x_j_framej_2 = x_j_framej.permute(0, 2, 3, 1) # (batch_size, num_points, k, num_dims)
+ pre = x_j_framej_2.reshape(-1, x_j_framej_2.shape[-1]) # (batch_size*num_points*k, num_dims)
+ x_j_framei = trafo(pre, trafo_j_to_i)
+ x_j_framei = x_j_framei.view(x_j_framej_2.shape).permute(
+ 0, 3, 1, 2
+ ) # (batch_size, num_dims, num_points, k)
+ return x_j_framei
def knn(x, k):
inner = -2 * torch.matmul(x.transpose(2, 1), x)
xx = torch.sum(x**2, dim=1, keepdim=True)
pairwise_distance = -xx - inner - xx.transpose(2, 1)
idx = pairwise_distance.topk(k=k + 1, dim=-1)[1][:, :, 1:] # (batch_size, num_points, k)
return idx
# v1 is faster on GPU
-def get_graph_feature_v1(x, k, idx):
+def get_graph_feature_v1(x, k, idx, frames, trafo):
batch_size, num_dims, num_points = x.size()
idx_base = torch.arange(0, batch_size, device=x.device).view(-1, 1, 1) * num_points
idx = idx + idx_base
idx = idx.view(-1)
fts = x.transpose(2, 1).reshape(
-1, num_dims
) # -> (batch_size, num_points, num_dims) -> (batch_size*num_points, num_dims)
fts = fts[idx, :].view(
batch_size, num_points, k, num_dims
) # neighbors: -> (batch_size*num_points*k, num_dims) -> ...
fts = fts.permute(0, 3, 1, 2).contiguous() # (batch_size, num_dims, num_points, k)
x = x.view(batch_size, num_dims, num_points, 1).repeat(1, 1, 1, k)
+ fts = change_local_frame(fts, idx, frames, trafo)
fts = torch.cat((x, fts - x), dim=1) # ->(batch_size, 2*num_dims, num_points, k)
return fts
# v2 is faster on CPU
-def get_graph_feature_v2(x, k, idx):
+def get_graph_feature_v2(x, k, idx, frames, trafo):
batch_size, num_dims, num_points = x.size()
idx_base = torch.arange(0, batch_size, device=x.device).view(-1, 1, 1) * num_points
idx = idx + idx_base
idx = idx.view(-1)
fts = x.transpose(0, 1).reshape(
num_dims, -1
) # -> (num_dims, batch_size, num_points) -> (num_dims, batch_size*num_points)
fts = fts[:, idx].view(
num_dims, batch_size, num_points, k
) # neighbors: -> (num_dims, batch_size*num_points*k) -> ...
fts = fts.transpose(1, 0).contiguous() # (batch_size, num_dims, num_points, k)
+ fts = change_local_frame(fts, idx, frames, trafo)
x = x.view(batch_size, num_dims, num_points, 1).repeat(1, 1, 1, k)
fts = torch.cat((x, fts - x), dim=1) # ->(batch_size, 2*num_dims, num_points, k)
return fts
class EdgeConvBlock(nn.Module):
r"""EdgeConv layer.
Introduced in "`Dynamic Graph CNN for Learning on Point Clouds
<https://arxiv.org/pdf/1801.07829>`__". Can be described as follows:
.. math::
x_i^{(l+1)} = \max_{j \in \mathcal{N}(i)} \mathrm{ReLU}(
\Theta \cdot (x_j^{(l)} - x_i^{(l)}) + \Phi \cdot x_i^{(l)})
where :math:`\mathcal{N}(i)` is the neighbor of :math:`i`.
Parameters
----------
in_feat : int
Input feature size.
out_feat : int
Output feature size.
batch_norm : bool
Whether to include batch normalization on messages.
"""
def __init__(
self,
k,
- in_feats
+ in_reps,
out_feats,
batch_norm=True,
activation=True,
cpu_mode=False,
):
super(EdgeConvBlock, self).__init__()
self.k = k
self.batch_norm = batch_norm
self.activation = activation
self.num_layers = len(out_feats)
self.get_graph_feature = get_graph_feature_v2 if cpu_mode else get_graph_feature_v1
+ in_feat = in_reps.dim
+ self.trafo = TensorRepsTransform(TensorReps(in_reps))
self.convs = nn.ModuleList()
for i in range(self.num_layers):
self.convs.append(
nn.Conv2d(
2 * in_feat if i == 0 else out_feats[i - 1],
out_feats[i],
kernel_size=1,
bias=False if self.batch_norm else True,
)
)
if batch_norm:
self.bns = nn.ModuleList()
for i in range(self.num_layers):
self.bns.append(nn.BatchNorm2d(out_feats[i]))
if activation:
self.acts = nn.ModuleList()
for i in range(self.num_layers):
self.acts.append(nn.ReLU())
if in_feat == out_feats[-1]:
self.sc = None
else:
self.sc = nn.Conv1d(in_feat, out_feats[-1], kernel_size=1, bias=False)
self.sc_bn = nn.BatchNorm1d(out_feats[-1])
if activation:
self.sc_act = nn.ReLU()
- def forward(self, points, features):
+ def forward(self, points, features, frames):
topk_indices = knn(points, self.k)
- x = self.get_graph_feature(features, self.k, topk_indices)
+ x = self.get_graph_feature(features, self.k, topk_indices, frames, self.trafo)
for conv, bn, act in zip(self.convs, self.bns, self.acts, strict=False):
x = conv(x) # (N, C', P, K)
if bn:
x = bn(x)
if act:
x = act(x)
fts = x.mean(dim=-1) # (N, C, P)
# shortcut
if self.sc:
sc = self.sc(features) # (N, C_out, P)
sc = self.sc_bn(sc)
else:
sc = features
return self.sc_act(sc + fts) # (N, C_out, P)
class ParticleNet(nn.Module):
"""ParticleNet with local frame transformations."""
def __init__(
self,
input_dims,
+ hidden_reps_list,
num_classes,
conv_params=[(7, (32, 32, 32)), (7, (64, 64, 64))],
fc_params=[(128, 0.1)],
use_fusion=True,
use_fts_bn=True,
use_counts=True,
for_inference=False,
for_segmentation=False,
**kwargs,
):
+ # hidden_reps_list: hidden representation for message-passing at beginning of each layer
super(ParticleNet, self).__init__(**kwargs)
+ hidden_reps_list = [TensorReps(x) for x in hidden_reps_list]
+ assert input_dims == hidden_reps_list[0].dim
+ assert len(hidden_reps_list) == len(conv_params)
self.use_fts_bn = use_fts_bn
if self.use_fts_bn:
self.bn_fts = nn.BatchNorm1d(hidden_reps_list[0].dim)
self.use_counts = use_counts
self.edge_convs = nn.ModuleList()
for idx, layer_param in enumerate(conv_params):
k, channels = layer_param
- in_feat = input_dims if idx == 0 else conv_params[idx - 1][1][-1]
+ in_reps = hidden_reps_list[idx]
+ assert (
+ in_reps.dim == conv_params[idx - 1][1][-1] if idx > 0 else hidden_reps_list[0].dim
+ )
self.edge_convs.append(
- EdgeConvBlock(k=k, in_feat=in_feat, out_feats=channels, cpu_mode=for_inference)
+ EdgeConvBlock(k=k, in_reps=in_reps, out_feats=channels, cpu_mode=for_inference)
)
self.use_fusion = use_fusion
if self.use_fusion:
in_chn = sum(x[-1] for _, x in conv_params)
out_chn = max(128, min((in_chn // 128) * 128, 1024))
self.fusion_block = nn.Sequential(
nn.Conv1d(in_chn, out_chn, kernel_size=1, bias=False),
nn.BatchNorm1d(out_chn),
nn.ReLU(),
)
self.for_segmentation = for_segmentation
fcs = []
for idx, layer_param in enumerate(fc_params):
channels, drop_rate = layer_param
if idx == 0:
in_chn = out_chn if self.use_fusion else conv_params[-1][1][-1]
else:
in_chn = fc_params[idx - 1][0]
if self.for_segmentation:
fcs.append(
nn.Sequential(
nn.Conv1d(in_chn, channels, kernel_size=1, bias=False),
nn.BatchNorm1d(channels),
nn.ReLU(),
nn.Dropout(drop_rate),
)
)
else:
fcs.append(
nn.Sequential(nn.Linear(in_chn, channels), nn.ReLU(), nn.Dropout(drop_rate))
)
if self.for_segmentation:
fcs.append(nn.Conv1d(fc_params[-1][0], num_classes, kernel_size=1))
else:
fcs.append(nn.Linear(fc_params[-1][0], num_classes))
self.fc = nn.Sequential(*fcs)
self.for_inference = for_inference
- def forward(self, points, features, mask=None):
+ def forward(self, points, features, frames, mask=None):
# print('points:\n', points)
# print('features:\n', features)
if mask is None:
mask = features.abs().sum(dim=1, keepdim=True) != 0 # (N, 1, P)
points *= mask
features *= mask
coord_shift = (mask == 0) * 1e9
if self.use_counts:
counts = mask.float().sum(dim=-1)
counts = torch.max(counts, torch.ones_like(counts)) # >=1
if self.use_fts_bn:
fts = self.bn_fts(features) * mask
else:
fts = features
outputs = []
for idx, conv in enumerate(self.edge_convs):
pts = (points if idx == 0 else fts) + coord_shift
- fts = conv(pts, fts) * mask
+ fts = conv(pts, fts, frames) * mask
if self.use_fusion:
outputs.append(fts)
if self.use_fusion:
fts = self.fusion_block(torch.cat(outputs, dim=1)) * mask
# assert(((fts.abs().sum(dim=1, keepdim=True) != 0).float() - mask.float()).abs().sum().item() == 0)
if self.for_segmentation:
x = fts
else:
if self.use_counts:
x = fts.sum(dim=-1) / counts # divide by the real counts
else:
x = fts.mean(dim=-1)
output = self.fc(x)
if self.for_inference:
output = torch.softmax(output, dim=1)
# print('output:\n', output)
return output