I have this code from DGCNN Neural Network but i don't understand how it extracts features. In particular i understand that we get the top knn point but i don't understand the idx_base.

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, dim=-1)[1]   # (batch_size, num_points, k)
  return idx

def get_graph_feature(x, k=20, idx=None):
  batch_size = x.size(0)
  num_points = x.size(2)
  x = x.view(batch_size, -1, num_points)
  if idx is None:
      idx = knn(x, k=k)   # (batch_size, num_points, k)
  device = torch.device('cuda')

  idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1)*num_points

  idx = idx + idx_base

  idx = idx.view(-1)

  _, num_dims, _ = x.size()

  x = x.transpose(2, 1).contiguous()   # (batch_size, num_points, num_dims)  -> (batch_size*num_points, num_dims) #   batch_size * num_points * k + range(0, batch_size*num_points)
  feature = x.view(batch_size*num_points, -1)[idx, :]
  feature = feature.view(batch_size, num_points, k, num_dims) 
  x = x.view(batch_size, num_points, 1, num_dims).repeat(1, 1, k, 1)

  feature = torch.cat((feature-x, x), dim=3).permute(0, 3, 1, 2).contiguous()

  return feature

1 Answer 1


i don't understand how it extracts features.

What you show is not really extracting any features, but collecting them:

  1. Compute a distance from each point to each other point in x (idx = knn(x, k))
  2. Collect the features of the k-nearest points = x.view(batch_size*num_points, -1)[idx, :]
  3. Now for each x you have k other vectors, so x needs to be repeated k-times: x.repeat(1, 1, k, 1)
  4. Compute the difference between x and all k nearest points and concat x: feature = torch.cat((feature-x, x), dim=3)

The last step provides you with one feature for each of the k-nearest neighbor points. Your features have now the shape: [ batch_size, num_dims, num_points, k ]. You can now apply some Neural Network to them.

[...] i don't understand the idx_base.

Something like this is common in GNNs and relates to graph batching. Notice the transpose:

(batch_size, num_points, num_dims) -> (batch_size*num_points, num_dims)

To get an index that works with this shape, you cannot use the raw node indices, you have to offset them using the batch number and the number of points per sample. This is what idx_base is doing, it offsets the point indices. This DGL-documentation has a nice visualization for what graph-batching means. The same applies here for idx_base.


Not the answer you're looking for? Browse other questions tagged .