Export densepose model to torchscript

[zimratK]

Hello,
I am trying to export a densepose model to TorchScript format for deployment. From this tutorial I understand that the common models can be converted to TorchScript format by tracing or scripting, but densepose may not be one of the common models. I tried to export the model using this code:

import sys
sys.path.append("../detectron2")
from detectron2.config import get_cfg
from detectron2.modeling import build_model
from detectron2.engine import DefaultPredictor
from detectron2.export import Caffe2Tracer, add_export_config
sys.path.append("../detectron2/projects/DensePose/")
from densepose import add_densepose_config
import cv2
import torch
import subprocess
cfg_path = "../detectron2/projects/DensePose/configs/densepose_rcnn_R_101_FPN_DL_s1x.yaml"
cfg = get_cfg()
add_densepose_config(cfg)
cfg.merge_from_file(cfg_path)
cfg.MODEL.DEVICE = "cpu"
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
cfg.MODEL.WEIGHTS = "../server/models/densepose/model_final_844d15.pkl"
add_export_config(cfg)
model = build_model(cfg)
model.eval()


#prepare model input
predictor = DefaultPredictor(cfg)
path_to_image = "/home/zimrat/Pictures/example.jpg"

original_image = cv2.imread(path_to_image)
with torch.no_grad():
    # Apply pre-processing to image.
    if predictor.input_format == "RGB":
        # whether the model expects BGR inputs or RGB
        original_image = original_image[:, :, ::-1]
    height, width = original_image.shape[:2]
    image = predictor.aug.get_transform(original_image).apply_image(original_image)
    image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))

    inputs = [{"image": image, "height": height, "width": width}]
caffe2_tracer = Caffe2Tracer(cfg, model, inputs)
caffe2_tracer.export_torchscript()

But it did not work. Can you give me a hint about how can I export the model?

The error I got:

warnings.warn(
/home/zimrat/projects/measure-kid-server/detectron2/detectron2/export/c10.py:31: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
assert tensor.dim() == 2 and tensor.size(-1) in [4, 5, 6], tensor.size()
/home/zimrat/projects/measure-kid-server/detectron2/detectron2/export/c10.py:370: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
assert box_regression.shape[1] % box_dim == 0
/home/zimrat/projects/measure-kid-server/detectron2/detectron2/export/c10.py:377: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
if input_tensor_mode:
/home/zimrat/projects/measure-kid-server/detectron2/detectron2/export/c10.py:409: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
nms_outputs = torch.ops._caffe2.BoxWithNMSLimit(
/home/zimrat/.venv/lib/python3.8/site-packages/torch/tensor.py:587: RuntimeWarning: Iterating over a tensor might cause the trace to be incorrect. Passing a tensor of different shape won't change the number of iterations executed (and might lead to errors or silently give incorrect results).
warnings.warn('Iterating over a tensor might cause the trace to be incorrect. '
/home/zimrat/projects/measure-kid-server/detectron2/detectron2/export/c10.py:438: TracerWarning: Converting a tensor to a Python number might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
for i, b in enumerate(int(x.item()) for x in roi_batch_splits_nms)
/home/zimrat/projects/measure-kid-server/SageMakerTry/../detectron2/projects/DensePose/densepose/data/structures.py:291: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
assert S.size()[2:] == I.size()[2:], (
/home/zimrat/projects/measure-kid-server/SageMakerTry/../detectron2/projects/DensePose/densepose/data/structures.py:296: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
assert I.size() == U.size(), (
/home/zimrat/projects/measure-kid-server/SageMakerTry/../detectron2/projects/DensePose/densepose/data/structures.py:300: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
assert I.size() == V.size(), (
Traceback (most recent call last):
File "export_densepose_model.py", line 52, in
caffe2_tracer.export_torchscript()
File "/home/zimrat/projects/measure-kid-server/detectron2/detectron2/export/api.py", line 140, in export_torchscript
return torch.jit.trace(model, (inputs,), optimize=True)
File "/home/zimrat/.venv/lib/python3.8/site-packages/torch/jit/_trace.py", line 733, in trace
return trace_module(
File "/home/zimrat/.venv/lib/python3.8/site-packages/torch/jit/_trace.py", line 934, in trace_module
module._c._create_method_from_trace(
RuntimeError: Tracer cannot infer type of (tensor([[0., 0., 0., 0.]]), tensor([1.]), tensor([0.]), <densepose.data.structures.DensePoseOutput object at 0x7f792950e520>)
:Only tensors and (possibly nested) tuples of tensors, lists, or dictsare supported as inputs or outputs of traced functions, but instead got value of type DensePoseOutput.

Thanks in advance!