C++ Inputs format for torchscript model generated via scripting export-method

[rds-itga]

Hi,
I am very new to torchscript and deployement and i would be very grateful if someone could help me.

I generated a torchscript model (model.ts) from a pre-trained mask-rcnn model.pt thanks to the python script export_model.py using the torchscript format and scripting export-method.
Then I tried to perform inference in C++ using torchscript_traced_mask_rcnn.cpp, i know that this code is dedicated to model.ts generated with tracing export-method, but i just wanted to see what would happened with a model.ts generated with scripting export-method.
As expected it did not work and i got this message:

terminate called after throwing an instance of 'c10::Error'
  what():  forward() Expected a value of type 'Tuple[Dict[str, Tensor]]' for argument 'inputs' but instead found type 'Tensor'.

So the forward method seems to expect a Tuple[Dict[str, Tensor]] as inputs but I don't know how to do it in C++.
So my question is what should be the get_scripting_inputs(cv::Mat& img, c10::Device device) function to make torchscript_traced_mask_rcnn.cpp working with a model.ts generated with the scripting export-method?

My second question is about batch size. I've read that the batch size is fixed for model.ts generated via tracing export-method.
For the moment export_model.py only permits a batch size of 1. I tried to modify the code to allow a batch size of 4 but it was unsuccessful.
What should i modify in

def export_tracing(torch_model, inputs):
    assert TORCH_VERSION >= (1, 8)
    image = inputs[0]["image"]
    inputs = [{"image": image}]  # remove other unused keys

    if isinstance(torch_model, GeneralizedRCNN):

        def inference(model, inputs):
            # use do_postprocess=False so it returns ROI mask
            inst = model.inference(inputs, do_postprocess=False)[0]
            return [{"instances": inst}]

    else:
        inference = None  # assume that we just call the model directly

    traceable_model = TracingAdapter(torch_model, inputs, inference)

    if args.format == "torchscript":
        ts_model = torch.jit.trace(traceable_model, (image,))
        with PathManager.open(os.path.join(args.output, "model.ts"), "wb") as f:
            torch.jit.save(ts_model, f)
        dump_torchscript_IR(ts_model, args.output)
    elif args.format == "onnx":
        # NOTE onnx export currently failing in pytorch
        with PathManager.open(os.path.join(args.output, "model.onnx"), "wb") as f:
            torch.onnx.export(traceable_model, (image,), f)
    logger.info("Inputs schema: " + str(traceable_model.inputs_schema))
    logger.info("Outputs schema: " + str(traceable_model.outputs_schema))

    if args.format != "torchscript":
        return None
    if not isinstance(torch_model, (GeneralizedRCNN, RetinaNet)):
        return None

to set a batch size of 4 for example?

Thanks in advance for your help.

[ppwwyyxx]

A reference is given in #2582 (comment)

[rds-itga]

Hi,
thank you very much for your answer, it is very helpful.

[rds-itga]

Concerning the definition of the batch size during export tracing I modified the python function export_tracing(torch_model, inputs).
For a batch size of 4 and assuming inputs variable being a list of 4 dicts, it gives:

def export_tracing(torch_model, inputs):
    assert TORCH_VERSION >= (1, 8)
    # image = inputs[0]["image"]
    # inputs = [{"image": image}]  # remove other unused keys
    image_0 = inputs[0]["image"]
    image_1 = inputs[1]["image"]
    image_2 = inputs[2]["image"]
    image_3 = inputs[3]["image"]
    inputs = [{"image": image_0}, {"image": image_1}, {"image": image_2}, {"image": image_3}]  # remove other unused keys

    if isinstance(torch_model, GeneralizedRCNN):

        def inference(model, inputs):
            # use do_postprocess=False so it returns ROI mask
            inst = model.inference(inputs, do_postprocess=False)[0]
            # return [{"instances": inst}]
            return [{"instances": inst[0]}, {"instances": inst[1]}, {"instances": inst[2]}, {"instances": inst[3]}]

    else:
        inference = None  # assume that we just call the model directly

    traceable_model = TracingAdapter(torch_model, inputs, inference)

    if args.format == "torchscript":
        # ts_model = torch.jit.trace(traceable_model, (image,))
        ts_model = torch.jit.trace(traceable_model, (image_0, image_1, image_2, image_3,))
        with PathManager.open(os.path.join(args.output, "model.ts"), "wb") as f:
            torch.jit.save(ts_model, f)
        dump_torchscript_IR(ts_model, args.output)
    elif args.format == "onnx":
        # NOTE onnx export currently failing in pytorch
        with PathManager.open(os.path.join(args.output, "model.onnx"), "wb") as f:
            # torch.onnx.export(traceable_model, (image,), f)
            torch.onnx.export(traceable_model, (image_0, image_1, image_2, image_3,), f)
    logger.info("Inputs schema: " + str(traceable_model.inputs_schema))
    logger.info("Outputs schema: " + str(traceable_model.outputs_schema))

    if args.format != "torchscript":
        return None
    if not isinstance(torch_model, (GeneralizedRCNN, RetinaNet)):
        return None

    def eval_wrapper(inputs):
        """
        The exported model does not contain the final resize step, which is typically
        unused in deployment but needed for evaluation. We add it manually here.
        """
        input = inputs[0]
        instances = traceable_model.outputs_schema(ts_model(input["image"]))[0]["instances"]
        postprocessed = detector_postprocess(instances, input["height"], input["width"])
        return [{"instances": postprocessed}]

    return eval_wrapper

then in C++ code I replaced module.forward({inputs}) by module.forward({inputs_0, inputs_1, inputs_2, inputs_3}). It seems to work but I don't know if it is the right way to do this or if there is more efficient way to do this?

Regards

[1695419]

Hb