[MRG][features] support onnx torchscript model conversion and flask server

model/model_deployor/deployor.py

@@ -67,22 +67,31 @@ def deploy(model,
         Reference:
             https://github.com/open-mmlab/mmdeploy/blob/master/tools/deploy.py
     """
-    assert backend in ['onnxruntime', 'tensorrt'], 'This backend isn\'t supported now!'
+    assert backend in ['onnxruntime', 'tensorrt', 'torchscript'], 'This backend isn\'t supported now!'
 
-    output_file = output_file + '.onnx'
-    torch.onnx.export(
-        model,
-        model_inputs,
-        output_file,
-        export_params=True,
-        input_names=input_names,
-        output_names=output_names,
-        opset_version=11,
-        dynamic_axes=dynamic_axes,
-        keep_initializers_as_inputs=False,
-        verbose=verbose)
     if backend == 'onnxruntime':
+        output_file = output_file + '.onnx'
+        torch.onnx.export(
+            model,
+            model_inputs,
+            output_file,
+            export_params=True,
+            input_names=input_names,
+            output_names=output_names,
+            opset_version=11,
+            dynamic_axes=dynamic_axes,
+            keep_initializers_as_inputs=False,
+            verbose=verbose)
         return output_file
+
+    if backend == 'torchscript':
+        with torch.no_grad():
+            output_file_name=output_file+".jit"
+            jit_model = torch.jit.trace(model, model_inputs)
+            jit_model.save(output_file_name)
+            print("torchscript successfully converts, saved as: ", output_file_name)
+            return output_file_name
+
     if backend == 'tensorrt':
         engine = create_trt_engine(
             output_file,

tools/deploy.py

@@ -10,23 +10,23 @@
 from mmcv.runner import load_checkpoint
 from model.model_deployor.deployor import deploy
 from model.model_deployor.deployor_utils import create_input
+import onnx
+import time
 import importlib
 if importlib.util.find_spec('tensorrt') is not None:
     from model.model_deployor.onnx2tensorrt import load_trt_engine, torch_dtype_from_trt, torch_device_from_trt
 else:
     print('Please install TensorRT if you want to convert')
 
-
 def main():
     parser = ArgumentParser()
     parser.add_argument('pcd', help='Point cloud file')
     parser.add_argument('checkpoint', help='Checkpoint file')
-    parser.add_argument('backend', default='onnx', help='backend name')
-    parser.add_argument('output', default='onnx', help='backend name')
-    parser.add_argument('dataset', default='onnx', help='backend name')
-    parser.add_argument('model_name', default='onnx', help='backend name')
-    parser.add_argument(
-        '--device', default='cuda:0', help='Device used for inference')
+    parser.add_argument('backend', default='onnx', help='support: onnxruntime, torchscript, tensorrt')
+    parser.add_argument('output', default='onnx', help='output model file name')
+    parser.add_argument('dataset', default='onnx', help='support: kitti, nuscenes')
+    parser.add_argument('model_name', default='onnx', help='support: pointpillars, centerpoint')
+    parser.add_argument('--device', default='cuda:0', help='Device used for inference')
 
     args = parser.parse_args()
 
@@ -37,6 +37,7 @@ def main():
         model = Centerpoint()
     load_checkpoint(model, args.checkpoint, map_location='cpu')
     model.cuda()
+    #model.cpu()
     model.eval()
 
     # define deploy params
@@ -56,6 +57,7 @@ def main():
 
     # verify
     torch_out = model(model_inputs[0], model_inputs[1], model_inputs[2])
+
     if args.backend == 'onnxruntime':
         import onnxruntime
 
@@ -71,9 +73,17 @@ def main():
         outputs['scores'] = torch.tensor(ort_output[0])
         outputs['bbox_preds'] = torch.tensor(ort_output[1])
         outputs['dir_scores'] = torch.tensor(ort_output[2])
-
         print('inference successful!')
 
+    if args.backend == 'torchscript':
+        jit_model=torch.jit.load(backend_file)
+        script_output=jit_model(model_inputs[0], model_inputs[1], model_inputs[2])
+        outputs = {}
+        outputs['scores'] = torch.tensor( [item.cpu().detach().numpy() for item in script_output[0]] )
+        outputs['bbox_preds'] = torch.tensor( [item.cpu().detach().numpy() for item in script_output[0]] )
+        outputs['dir_scores'] = torch.tensor( [item.cpu().detach().numpy() for item in script_output[0]] )
+        print("torchscript inference successful")
+
     if args.backend == 'tensorrt':
         engine = load_trt_engine(backend_file)
         context = engine.create_execution_context()
@@ -127,6 +137,5 @@ def main():
 
 
 
-
 if __name__ == '__main__':
     main()

tools/server.py

@@ -0,0 +1,212 @@
+# coding=utf-8
+
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = "3"
+
+import torch
+from argparse import ArgumentParser
+
+from deephub.detection_model import Pointpillars, Centerpoint
+
+from mmcv.runner import load_checkpoint
+from model.model_deployor.deployor import deploy
+from model.model_deployor.deployor_utils import create_input
+import onnx
+import onnxruntime
+import time
+import importlib
+if importlib.util.find_spec('tensorrt') is not None:
+    from model.model_deployor.onnx2tensorrt import load_trt_engine, torch_dtype_from_trt, torch_device_from_trt
+else:
+    print('Please install TensorRT if you want to convert')
+
+# --------------------------flask------------------------------------------
+
+import io
+import json
+
+import flask
+import torch
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from torch import nn
+from torchvision import transforms as T
+from torchvision.models import resnet50
+
+# Initialize our Flask application and the PyTorch model.
+app = flask.Flask(__name__)
+model = None
+use_gpu = True
+
+
+@app.route("/transfer", methods=["POST"])
+def transfer():
+    # Initialize the data dictionary that will be returned from the view.
+    if flask.request.method == 'POST':
+
+        pcd=flask.request.files["pcd"].read().decode("utf-8")
+        checkpoint = flask.request.files["checkpoint"].read().decode("utf-8")
+        backend = flask.request.files["backend"].read().decode("utf-8")
+        output = flask.request.files["output"].read().decode("utf-8")
+        dataset = flask.request.files["dataset"].read().decode("utf-8")
+        model_type = flask.request.files["model_type"].read().decode("utf-8")
+        device = flask.request.files["device"].read().decode("utf-8")
+
+        # Init model and load checkpoints
+        if model_type == 'pointpillars':
+            model = Pointpillars()
+        elif model_type == 'centerpoint':
+            model = Centerpoint()
+
+        load_checkpoint(model, checkpoint, map_location='cpu')
+        model.cuda()
+        # model.cpu()
+        model.eval()
+
+        # define deploy params
+        input_names = ['voxels', 'num_points', 'coors']
+        output_names = ['scores', 'bbox_preds', 'dir_scores']
+        dynamic_axes = {'voxels': {0: 'voxels_num'},
+                        'num_points': {0: 'voxels_num'},
+                        'coors': {0: 'voxels_num'}}
+        # dynamic_axes = None
+        fp16 = False
+
+        data, model_inputs = create_input(pcd, dataset, model_type, device)
+
+        # deploy
+        backend_file = deploy(model, model_inputs, input_names, output_names, dynamic_axes,
+                              backend=backend, output_file=output, fp16=fp16, dataset=dataset)
+
+        # verify
+        torch_out = model(model_inputs[0], model_inputs[1], model_inputs[2])
+
+        result={"status" : False}
+
+        if backend == 'onnxruntime':
+
+            ort_session = onnxruntime.InferenceSession(backend_file)
+
+            input_dict = {}
+            input_dict['voxels'] = model_inputs[0].cpu().numpy()
+            input_dict['num_points'] = model_inputs[1].cpu().numpy()
+            input_dict['coors'] = model_inputs[2].cpu().numpy()
+            ort_output = ort_session.run(['scores', 'bbox_preds', 'dir_scores'], input_dict)
+
+            outputs = {}
+            outputs['scores'] = torch.tensor(ort_output[0])
+            outputs['bbox_preds'] = torch.tensor(ort_output[1])
+            outputs['dir_scores'] = torch.tensor(ort_output[2])
+            print('inference successful!')
+
+            result["status"] = True
+            result["model_path"] = output+backend
+
+            return flask.jsonify(result)
+
+        if backend == 'torchscript':
+            jit_model = torch.jit.load(backend_file)
+            script_output = jit_model(model_inputs[0], model_inputs[1], model_inputs[2])
+            outputs = {}
+            outputs['scores'] = torch.tensor([item.cpu().detach().numpy() for item in script_output[0]])
+            outputs['bbox_preds'] = torch.tensor([item.cpu().detach().numpy() for item in script_output[0]])
+            outputs['dir_scores'] = torch.tensor([item.cpu().detach().numpy() for item in script_output[0]])
+
+            print("torchscript inference successful")
+
+            result["status"] = True
+            result["model_path"] = output+backend
+
+            return flask.jsonify(result)
+
+
+@app.route("/predict", methods=["POST"])
+def predict():
+    # Initialize the data dictionary that will be returned from the view.
+    if flask.request.method == 'POST':
+
+        pcd=flask.request.files["pcd"].read().decode("utf-8")
+        checkpoint = flask.request.files["checkpoint"].read().decode("utf-8")
+        backend = flask.request.files["backend"].read().decode("utf-8")
+        output = flask.request.files["output"].read().decode("utf-8")
+        dataset = flask.request.files["dataset"].read().decode("utf-8")
+        model_type = flask.request.files["model_type"].read().decode("utf-8")
+        device = flask.request.files["device"].read().decode("utf-8")
+
+        # Init model and load checkpoints
+        if model_type == 'pointpillars':
+            model = Pointpillars()
+        elif model_type == 'centerpoint':
+            model = Centerpoint()
+
+        load_checkpoint(model, checkpoint, map_location='cpu')
+        model.cuda()
+        # model.cpu()
+        model.eval()
+
+        # define deploy params
+        input_names = ['voxels', 'num_points', 'coors']
+        output_names = ['scores', 'bbox_preds', 'dir_scores']
+        dynamic_axes = {'voxels': {0: 'voxels_num'},
+                        'num_points': {0: 'voxels_num'},
+                        'coors': {0: 'voxels_num'}}
+        # dynamic_axes = None
+        fp16 = False
+
+        data, model_inputs = create_input(pcd, dataset, model_type, device)
+
+        # verify
+        torch_out = model(model_inputs[0], model_inputs[1], model_inputs[2])
+
+        result={"status" : False}
+
+        if backend == 'onnxruntime':
+
+            backend_file=output+".onnx"
+
+            ort_session = onnxruntime.InferenceSession(backend_file)
+
+            input_dict = {}
+            input_dict['voxels'] = model_inputs[0].cpu().numpy()
+            input_dict['num_points'] = model_inputs[1].cpu().numpy()
+            input_dict['coors'] = model_inputs[2].cpu().numpy()
+            ort_output = ort_session.run(['scores', 'bbox_preds', 'dir_scores'], input_dict)
+
+            outputs = {}
+            outputs['scores'] = torch.tensor(ort_output[0])
+            outputs['bbox_preds'] = torch.tensor(ort_output[1])
+            outputs['dir_scores'] = torch.tensor(ort_output[2])
+            print('inference successful!')
+
+            result["status"] = True
+            result['scores'] = ort_output[0].tolist()
+            result['bbox_preds'] = ort_output[1].tolist()
+            result['coors'] = ort_output[2].tolist()
+
+            return flask.jsonify(result)
+
+        if backend == 'torchscript':
+
+            backend_file=output+".jit"
+            jit_model = torch.jit.load(backend_file)
+            script_output = jit_model(model_inputs[0], model_inputs[1], model_inputs[2])
+            outputs = {}
+            outputs['scores'] = torch.tensor([item.cpu().detach().numpy() for item in script_output[0]])
+            outputs['bbox_preds'] = torch.tensor([item.cpu().detach().numpy() for item in script_output[0]])
+            outputs['dir_scores'] = torch.tensor([item.cpu().detach().numpy() for item in script_output[0]])
+
+            print("torchscript inference successful")
+
+            result["status"] = True
+            result['scores'] = ort_output[0].tolist()
+            result['bbox_preds'] = ort_output[1].tolist()
+            result['coors'] = ort_output[2].tolist()
+
+            return flask.jsonify(result)
+
+
+
+if __name__ == '__main__':
+    print("Loading PyTorch model and Flask starting server ... \nPlease wait until server has fully started")
+    app.run(port=1234)