GraspNet-Edge is an optimized and efficient 6D grasp pose predictor based on GraspNet, designed for edge devices.
During the deployment of GraspNet to edge platforms, we observed that many operators were missing when exporting the model to ONNX format. After root cause analysis, we identified that the missing operators originated from CUDA-defined custom layers in PointNet++ (pointnet2).
To address this issue, we reimplemented these CUDA-based operators using native PyTorch operations and provided corresponding code to enable successful ONNX export.
In addition, we provide parallel training code based on the DDP (Distributed Data Parallel) framework.
Optimized GraspNet For Edge Devices !
Explore the documentation of this project »
Demo
·
Report Bug
·
Propose New Feature
- Wrote the code in the
pointnet2/pointnet2_modules_pytorch.pyfile, which contains operators implemented purely in PyTorch. - Modified the model code in
models/backbone.py,models/graspnet.pyandmodels/modules.py. Modifications to these files include replacing CUDA-defined operators with PyTorch-defined ones, as well as converting 1D convolutions into equivalent 2D convolutions to better adapt to the hardware characteristics of the Horizon X5 RDK edge device. - You can refer to the
train_distributed.pyfile to see how to use DDP to train the model.
python3 to_onnx.py
Hongrui Zhu
E-Mail:[email protected] or [email protected]
qq:786739982
vx:Hong_Rui_0226
该项目签署了MIT 授权许可,详情请参阅 LICENSE