本文介绍如何用 Relay VM 部署 PyTorch 目标检测模型。 首先应安装 PyTorch。此外，还应安装 TorchVision，并将其作为模型合集（model zoo）。 可通过 pip 快速安装： pip install torchpip install torchvision或参考官网：https://pytorch.org/get-started/locally/ PyTorch 版本应该和 TorchVision 版本兼容。 目前 TVM 支持 PyTorch 1.7 和 1.4，其他版本可能不稳定。 import tvmfrom tvm import relayfrom tvm import relayfrom tvm.runtime.vm import VirtualMachinefrom tvm.contrib.download import download_testdataimport numpy as npimport cv2# PyTorch 导入import torchimport torchvision从 TorchVision 加载预训练的 MaskRCNN 并进行跟踪in_size = 300input_shape = (1, 3, in_size, in_size)def do_trace(model, inp): model_trace = torch.jit.trace(model, inp) model_trace.eval() return model_tracedef 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_fpnmodel = 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)输出结果： Downloading: "https://download.pytorch.org/models/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth" to /workspace/.cache/torch/hub/checkpoints/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth 0%| | 0.00/170M [00:00 score_threshold: valid_boxes.append(boxes[i]) else: breakprint("Get {} valid boxes".format(len(valid_boxes)))输出结果： Get 9 valid boxes脚本总运行时长：（2 分 57.278 秒） 下载 Python 源代码：deploy_object_detection_pytorch.py 下载 Jupyter Notebook：deploy_object_detection_pytorch.ipynb