Self-supervised Graph Neural Network for Mechanical CAD Retrieval

TL;DR

GC-CAD tackles CAD similarity retrieval by learning structure-aware representations from BRep-derived graphs with a self-supervised contrastive framework. It combines a UVNet-based feature extractor with graph neural networks to capture both local geometric details and global topology, trained entirely without labels via a NT-Xent loss and graph augmentations. The approach delivers clear accuracy gains over baselines and offers up to 100× efficiency improvements through vector-based retrieval and FAISS indexing. Its validation on public and private datasets, including large-scale and assembly retrieval scenarios, demonstrates strong practical potential for scalable CAD reuse in industry.

Abstract

CAD (Computer-Aided Design) plays a crucial role in mechanical industry, where large numbers of similar-shaped CAD parts are often created. Efficiently reusing these parts is key to reducing design and production costs for enterprises. Retrieval systems are vital for achieving CAD reuse, but the complex shapes of CAD models are difficult to accurately describe using text or keywords, making traditional retrieval methods ineffective. While existing representation learning approaches have been developed for CAD, manually labeling similar samples in these methods is expensive. Additionally, CAD models' unique parameterized data structure presents challenges for applying existing 3D shape representation learning techniques directly. In this work, we propose GC-CAD, a self-supervised contrastive graph neural network-based method for mechanical CAD retrieval that directly models parameterized CAD raw files. GC-CAD consists of two key modules: structure-aware representation learning and contrastive graph learning framework. The method leverages graph neural networks to extract both geometric and topological information from CAD models, generating feature representations. We then introduce a simple yet effective contrastive graph learning framework approach, enabling the model to train without manual labels and generate retrieval-ready representations. Experimental results on four datasets including human evaluation demonstrate that the proposed method achieves significant accuracy improvements and up to 100 times efficiency improvement over the baseline methods.

Figures (8)

• Figure 1: Boundary representation(BRep) and converting it to graph. The faces and curves in BRep correspond to the nodes and edges in graph.
• Figure 2: Overall framework of GC-CAD. Each CAD part is first converted to a graph, then feature and structure augmentation are performed separately. Augmented graph is input to GNN to obtain graph representation. Finally, GC-CAD constructs and learns from positive and negative sample pairs based on whether the graph representation are from the same original CAD part. After the training is completed, the CAD retrieval can be performed using graph representation.
• Figure 3: Retrieved results of different data augmentation method.(a) mask node. (b) mask node and its 1-hop neighbors. (c) mask edge and its vertices.
• Figure 4: Retrieved results on large scale dataset.
• Figure 5: Retrieved results for generalization experiments.(a)Training on ABC dataste. (b)Training on ABC_small dataset.