MeshCNN: symmetric features for equivariant convolution

mesh_conv.py

# Simplified from models/layers/mesh_conv.py in https://github.com/ranahanocka/MeshCNN

class MeshCov(nn.Module):
  def __init__(self, in_c, out_c, k=5, bias=True):
    super(MeshConv, self).__init__()
    self.conv = nn.Conv2d(in_channels=in_c, out_channels=out_c, kernel_size=(1, k), bias=bias)
  
  def forward(self, x)
    """
      Forward pass given a feature tensor x with shape (N, C, E, 5):
        N - batch
        C - # features
        E - # edges in mesh
        5 - edges in neighborhood (0 is central edge)
    """
    x_1 = x[:, :, :, 1] + x[:, :, :, 3]
    x_2 = x[:, :, :, 2] + x[:, :, :, 4]
    x_3 = torch.abs(x[:, :, :, 1] - x[:, :, :, 3])
    x_4 = torch.abs(x[:, :, :, 2] - x[:, :, :, 4])
    x = torch.stack([x[:, :, :, 0], x_1, x_2, x_3, x_4], dim=3)
    x = self.conv(x)
    return x