martinnormark / neural-mesh-simplification

Implementation of the Neural Mesh Simplification paper
MIT License
3 stars 0 forks source link

Implement `face_classifier` for triangle probability #18

Closed martinnormark closed 1 month ago

martinnormark commented 1 month ago

Key components and breakdown:

  1. k-nn graph construction: We're using PyTorch Geometric's knn_graph function to construct the k-nn graph based on barycenter distances.
  2. TriConv layer: We need to implement the TriConv operation as described in the paper. This involves computing the relative position encoding and applying the update rule.
  3. Stacked TriConv layers: We're using a ModuleList to stack multiple TriConv layers.
  4. Final classification: We use a linear layer followed by a softmax activation to produce the final probabilities.
martinnormark commented 1 month ago

The FaceClassifier implementation compared to the description in the paper:

  1. Initialization:

    • The class takes input_dim, hidden_dim, num_layers, and k as parameters, which aligns with the paper's description of a configurable architecture.
    • It creates a ModuleList of TriConv layers, which corresponds to the paper's mention of stacked TriConv layers.
    • The final linear layer is used to produce the output, which is then transformed into probabilities.
  2. Forward method:

    • It computes barycenters if the input positions are 3D, which aligns with the paper's description of using face barycenters.
    • It constructs a k-nn graph using a custom method, which is in line with the paper's mention of connecting each candidate triangle with its k-nearest neighbors.
    • It applies the TriConv layers sequentially with ReLU activation, which matches the paper's description.
    • The final classification is done using a linear layer, followed by softmax (or scatter_softmax for batched inputs), which aligns with the paper's mention of using softmax to produce the final probabilities.
  3. custom_knn_graph method:

    • This method constructs a k-nearest neighbors graph, which is consistent with the paper's description.
    • It handles batched inputs, which is not explicitly mentioned in the paper but is a useful feature for efficient processing.
    • The graph is made symmetric, which is not explicitly mentioned in the paper but is a common practice in graph neural networks to ensure information flows both ways.

Overall, the implementation appears to be in line with the paper's description of the Face Classifier. The main components described in the paper (k-nn graph construction, stacked TriConv layers, and final classification) are all present.

However, there are a few points to consider:

  1. The paper mentions using barycenter distances for the k-nn graph, which is implicitly done here by using the mean of 3D positions. This could be made more explicit if needed.

  2. The custom_knn_graph method doesn't use PyTorch Geometric's knn_graph function. While this is fine given that the tests pass, using PyTorch Geometric's implementation could potentially be more efficient and better optimized.

  3. The paper doesn't specify the exact activation function used between TriConv layers. The implementation uses ReLU, which is a common choice and should work well.

  4. The implementation includes batch handling, which is not mentioned in the paper but is a practical addition for efficient processing of multiple meshes.

In conclusion, while there are minor differences and additions, the implementation appears to be faithful to the paper's description of the Face Classifier. The changes and additions (like batch handling and graph symmetry) are practical improvements that don't fundamentally alter the described architecture.