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

General Discussion
,
#1

I am using detectron2 model for instant segmentation & object detection .My goal is to deploy this model onto a robot using Viam robotics 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 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’) & convert it to tflite. I would like to get your help on this, please.

#3

@Kamal_Moha

In the context of Detectron2, the “–d2-pkl” flag in the provided script stands for the path to the serialized PyTorch model weights in pickle format. This file contains the learned parameters of the model after training.

Now, to convert your custom trained model in ‘.pth’ format to TensorFlow format (either ‘.pb’ or ‘.tflite’), you’ll need to follow a series of steps.

pip install torch torchvision onnx tf2onnx

import torch.onnx
import torchvision.models as models

# Load the PyTorch model
model = models.detection.maskrcnn_resnet50_fpn(pretrained=False)
model.load_state_dict(torch.load('path_to_your_custom_model.pth'))
model.eval()

# Dummy input (adjust according to your model's input specifications)
dummy_input = torch.randn(1, 3, 224, 224)

# Export to ONNX format
onnx_path = 'path_to_output_model.onnx'
torch.onnx.export(model, dummy_input, onnx_path, verbose=True, input_names=['input'], output_names=['output'])
  1. Convert ONNX to TensorFlow (saved model)
python -m tf2onnx.convert --opset 11 --fold_const -t tf --input path_to_output_model.onnx --output path_to_output_model.pb
  1. Convert to TensorFlow Lite (.tflite)
tflite_convert --output_file=path_to_output_model.tflite --saved_model_dir=path_to_saved_model_directory

You might also need to adjust the input size and format based on your custom model’s specifications. Also, the models.detection.maskrcnn_resnet50_fpn in the script is just an example; you should replace it with the architecture you used for training your custom model.

Once you have the ‘.pb’ file, you can then proceed to convert it to ‘.tflite’ if needed, as shown in the optional step.

#5

I’m getting errors when I try to convert ONNX to tf. I have managed to convert the pytorch(.pth) model to ONNX using an inbuilt way using export_model.py(made a few edits) in detectron2 using the code;

%run /content/export_model1.py --config-file /content/output.yaml --output /content/model.onnx --format onnx --sample-image /content/32700.jpg --export-method tracing MODEL.DEVICE cuda MODEL.WEIGHTS /content/output/model_final.pth

The above code was able to convert pytorch to ONNX, but can’t convert ONNX to tf

When I run the code;
!python -m tf2onnx.convert --opset 11 --fold_const -t tf --input /content/model.onnx/model.onnx --output /content/new_conv.pb

I get the error;
convert.py: error: unrecognized arguments: --fold_const -t tf

Even after removing --fold_const -t tf and run the code;
!python -m tf2onnx.convert --opset 11 --input /content/model.onnx/model.onnx --output /content/new_conv.pb I get the error;
convert.py: error: graphdef and checkpoint models need to provide inputs and outputs

I have googled for alternatives and found the library onnx-tf. I used onnx-tf and run the below code;

import onnx, os
from onnx_tf.backend import prepare
onnx_model = onnx.load('/content/model.onnx/model.onnx')   #load onnx model
tf_rep = prepare(onnx_model)            #prepare tf representation
tf_rep.export_graph("/content/new_conv.pb")   #export the model
assert os.path.exists(pb_path)
print(".pb model converted successfully.")

I get the error below after running it;

RuntimeError: in user code:

    File "/usr/local/lib/python3.10/dist-packages/onnx_tf/backend_tf_module.py", line 99, in __call__  *
        output_ops = self.backend._onnx_node_to_tensorflow_op(onnx_node,
    File "/usr/local/lib/python3.10/dist-packages/onnx_tf/backend.py", line 347, in _onnx_node_to_tensorflow_op  *
        return handler.handle(node, tensor_dict=tensor_dict, strict=strict)
    File "/usr/local/lib/python3.10/dist-packages/onnx_tf/handlers/handler.py", line 58, in handle  *
        cls.args_check(node, **kwargs)
    File "/usr/local/lib/python3.10/dist-packages/onnx_tf/handlers/backend/onehot.py", line 61, in args_check  *
        exception.DTYPE_NOT_CAST_EXCEPT(
    File "/usr/local/lib/python3.10/dist-packages/onnx_tf/common/exception.py", line 78, in __call__  *
        raise self._func(self.get_message(op, supported_dtypes))

RuntimeError: OneHot input onnx::OneHot_2278 with data type 'int64' is not supported in Tensorflow. Please set auto_cast to True or change data type to one of the supported types in the following list ['int32'].

I need help on how to solve this, please.