Kang2691196427/GRM

Structure of the Project

The embedding directory within the datasets folder stores embeddings of CAD parts obtained through pre-trained models. Specifically, embedding_comp2vec.json contains part embeddings generated using comp2vec, while embedding_uv.json contains part embeddings derived from UV-Net. The results.zip directory stores the trained model weights for 12 different methods.

---

Environment Setup

Create and activate the Conda environment:

conda env create -f environment.yml
conda activate CADRec

---

Training

• --GNN specifies the chosen graph neural network, which can be selected from GCN or GAT.
• --RM specifies the chosen recommendation model, which can be selected from dp (dot product), mlp, and ca (cross-attention).
• --embedding specifies the method for obtaining part embeddings, which can be selected from comp2vec and uv.
• --node_feature_dim specifies the dimension of the part embeddings: 128 for comp2vec, 64 for uv.
• --output_dim specifies the output dimension of the graph neural network, which is consistent with the part embedding dimension: 128 for comp2vec, 64 for uv.
• --experiment_name specifies the name of the experiment, used for storing checkpoints.
• --max_epochs specifies the maximum number of training epochs.

A code example for training the comp2vec + GCN + dp method is as follows:

python main.py train --GNN GCN --RM dp --embedding comp2vec --node_feature_dim 128 --output_dim 128 --experiment_name GCN_dp_comp2vec_450 --max_epochs 450

A code example for training the uv + GCN + ca method is as follows:

python main.py train --GNN GCN --RM ca --embedding uv --node_feature_dim 64 --output_dim 64 --experiment_name GCN_ca_uv_450 --max_epochs 450

---

Testing

• --GNN specifies the chosen graph neural network, which can be selected from GCN or GAT.
• --RM specifies the chosen recommendation model, which can be selected from dp (dot product), mlp, and ca (cross-attention).
• --embedding specifies the method for obtaining part embeddings, which can be selected from comp2vec and uv.
• --checkpoint specifies the storage location of the trained model checkpoints. The results folder contains the trained models for 12 different methods. Choose the best performing checkpoint from last.ckpt and best_loss.ckpt.

A code example for testing the comp2vec + GCN + dp method is as follows:

python main.py test --GNN GCN --RM dp --embedding comp2vec --checkpoint ./results/GCN_dp_comp2vec_450/0401/201405/last.ckpt

A code example for testing the uv + GCN + dp method is as follows:

python main.py test --GNN GCN --RM dp --embedding uv --checkpoint ./results/GCN_dp_uv_450/0403/163236/best-loss.ckpt