diff --git a/tutorials/frontend/deploy_object_detection_pytorch.py b/tutorials/frontend/deploy_object_detection_pytorch.py new file mode 100644 index 000000000000..6408685febfb --- /dev/null +++ b/tutorials/frontend/deploy_object_detection_pytorch.py @@ -0,0 +1,154 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Compile PyTorch Object Detection Models +======================================= +This article is an introductory tutorial to deploy PyTorch object +detection models with Relay VM. + +For us to begin with, PyTorch should be installed. +TorchVision is also required since we will be using it as our model zoo. + +A quick solution is to install via pip + +.. code-block:: bash + + pip install torch==1.4.0 + pip install torchvision==0.5.0 + +or please refer to official site +https://pytorch.org/get-started/locally/ + +PyTorch versions should be backwards compatible but should be used +with the proper TorchVision version. + +Currently, TVM supports PyTorch 1.4 and 1.3. Other versions may +be unstable. +""" + +import tvm +from tvm import relay +from tvm import relay +from tvm.runtime.vm import VirtualMachine +from tvm.contrib.download import download + +import numpy as np +import cv2 + +# PyTorch imports +import torch +import torchvision + +###################################################################### +# Load pre-trained maskrcnn from torchvision and do tracing +# --------------------------------------------------------- +in_size = 300 + +input_shape = (1, 3, in_size, in_size) + + +def do_trace(model, inp): + model_trace = torch.jit.trace(model, inp) + model_trace.eval() + return model_trace + + +def dict_to_tuple(out_dict): + if "masks" in out_dict.keys(): + return out_dict["boxes"], out_dict["scores"], out_dict["labels"], out_dict["masks"] + return out_dict["boxes"], out_dict["scores"], out_dict["labels"] + + +class TraceWrapper(torch.nn.Module): + def __init__(self, model): + super().__init__() + self.model = model + + def forward(self, inp): + out = self.model(inp) + return dict_to_tuple(out[0]) + + +model_func = torchvision.models.detection.maskrcnn_resnet50_fpn +model = TraceWrapper(model_func(pretrained=True)) + +model.eval() +inp = torch.Tensor(np.random.uniform(0.0, 250.0, size=(1, 3, in_size, in_size))) + +with torch.no_grad(): + out = model(inp) + script_module = do_trace(model, inp) + +###################################################################### +# Download a test image and pre-process +# ------------------------------------- +img_path = "test_street_small.jpg" +img_url = ( + "https://raw.githubusercontent.com/dmlc/web-data/" "master/gluoncv/detection/street_small.jpg" +) +download(img_url, img_path) + +img = cv2.imread(img_path).astype("float32") +img = cv2.resize(img, (in_size, in_size)) +img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) +img = np.transpose(img / 255.0, [2, 0, 1]) +img = np.expand_dims(img, axis=0) + +###################################################################### +# Import the graph to Relay +# ------------------------- +input_name = "input0" +shape_list = [(input_name, input_shape)] +mod, params = relay.frontend.from_pytorch(script_module, shape_list) + +###################################################################### +# Compile with Relay VM +# --------------------- +# Note: Currently only CPU target is supported. For x86 target, it is +# highly recommended to build TVM with Intel MKL and Intel OpenMP to get +# best performance, due to the existence of large dense operator in +# torchvision rcnn models. + +# Add "-libs=mkl" to get best performance on x86 target. +# For x86 machine supports AVX512, the complete target is +# "llvm -mcpu=skylake-avx512 -libs=mkl" +target = "llvm" + +with tvm.transform.PassContext(opt_level=3, disabled_pass=["FoldScaleAxis"]): + vm_exec = relay.vm.compile(mod, target=target, params=params) + +###################################################################### +# Inference with Relay VM +# ----------------------- +ctx = tvm.cpu() +vm = VirtualMachine(vm_exec, ctx) +vm.set_input("main", **{input_name: img}) +tvm_res = vm.run() + +###################################################################### +# Get boxes with score larger than 0.9 +# ------------------------------------ +score_threshold = 0.9 +boxes = tvm_res[0].asnumpy().tolist() +valid_boxes = [] +for i, score in enumerate(tvm_res[1].asnumpy().tolist()): + if score > score_threshold: + valid_boxes.append(boxes[i]) + else: + break + +print("Get {} valid boxes".format(len(valid_boxes)))