facebookresearch / facebookresearch/detectron2
How can I convert detecron2 (.pth) custom trained model to onnx format or .pb format(tflite) for robot deployment
- Dominant language
- Python
- Stars
- 34.7k
- Forks
- 7.9k
- PR merge metrics
- No merged PRs in 30d
Description
I am using detectron2 model for instant segmentation & object detection .My goal is to deploy this model onto a robot using [Viam robotics](https://www.viam.com/) and they(Viam) only accept a .tflite model when uploading.
I need a script (written in Python) that would take a trained PyTorch model file (.pth extension) and export it to TensorFlow format (.pb). [frozen graph] or .onnx so that I can be able to upload it into the robot
The model was trained using the Facebook's DETECTRON2 (the pre-trained model was "COCO-InstanceSegmentation/mask_rcnn_X_101_32x8d_FPN_3x.yaml")
I already have the output "pth" file. The script should take this file as a parameter and return single '.pb' file.
I checked this [doc](https://github.com/tensorpack/tensorpack/tree/master/examples/FasterRCNN/convert_d2) doing tensorflow conversions and it uses the below code for conversion;
```
# 1. Download the corresponding model from detectron2 model zoo
# 2. Convert:
$ python convert_d2.py --d2-config detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml --d2-pkl model_final_f10217.pkl --output R50FPN-d2-converted.npz
# the script will print tensorpack configs
'MODE_MASK=True' 'MODE_FPN=True' 'BACKBONE.STRIDE_1X1=True' 'PREPROC.PIXEL_MEAN=[123.675,116.28,103.53]' 'PREPROC.PIXEL_STD=[1.0,1.0,1.0]'
# 3. Use the above configs to verify the conversion is correct:
$ ./predict.py --evaluate out.json --load R50FPN-d2-converted.npz --config DATA.BASEDIR=~/data/coco 'MODE_MASK=True' 'MODE_FPN=True' 'BACKBONE.STRIDE_1X1=True' 'PREPROC.PIXEL_MEAN=[123.675,116.28,103.53]' 'PREPROC.PIXEL_STD=[1.0,1.0,1.0]'
# 4. Naively convert the model to a frozen pb file:
$ ./predict.py --output-pb out.pb --load R50FPN-d2-converted.npz --config DATA.BASEDIR=~/data/coco 'MODE_MASK=True' 'MODE_FPN=True' 'BACKBONE.STRIDE_1X1=True' 'PREPROC.PIXEL_MEAN=[123.675,116.28,103.53]' 'PREPROC.PIXEL_STD=[1.0,1.0,1.0]'
```
But I would like to know what `--d2-pkl` stands for and how can I use my custom trained model('.pth') for this conversion. I would like to get your help on this, please.
Contributor guide
Research direction
Start with the linked tensorpack examples/FasterRCNN/convert_d2/convert_d2.py workflow and the detectron2 model-zoo configuration named in the issue. Determine whether the custom .pth checkpoint can follow that conversion path and whether the resulting ONNX or frozen TensorFlow graph is accepted by Viam as TFLite; done means a verified conversion procedure for this custom Mask R-CNN model.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch, tensorflow
- Domain
- computer-vision, machine-learning, robotics
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 15/100