facebookresearch / facebookresearch/detectron2

How can I convert detecron2 (.pth) custom trained model to onnx format or .pb format(tflite) for robot deployment

Open
#5,180 1 comment 0 reactions 0 assignees View on GitHub
documentation
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.