[Fix] Fix the torchserve. (#1143)
* rebase * update docker and rm deprecated deployment tools * update docs * rebase Co-authored-by: Ezra-Yu <18586273+Ezra-Yu@users.noreply.github.com>pull/1211/head
marouane amzil
GitHub
parent
940a06f645
commit
743ca2d602
|
|
@ -4,6 +4,10 @@ ARG CUDNN="7"
|
|||
|
||||
FROM pytorch/pytorch:${PYTORCH}-cuda${CUDA}-cudnn${CUDNN}-devel
|
||||
|
||||
# fetch the key refer to https://forums.developer.nvidia.com/t/18-04-cuda-docker-image-is-broken/212892/9
|
||||
RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub 32
|
||||
RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/7fa2af80.pub
|
||||
|
||||
ENV TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0+PTX"
|
||||
ENV TORCH_NVCC_FLAGS="-Xfatbin -compress-all"
|
||||
ENV CMAKE_PREFIX_PATH="(dirname(which conda))/../"
|
||||
|
|
|
|||
|
|
@ -3,6 +3,10 @@ ARG CUDA="10.2"
|
|||
ARG CUDNN="7"
|
||||
FROM pytorch/pytorch:${PYTORCH}-cuda${CUDA}-cudnn${CUDNN}-devel
|
||||
|
||||
# fetch the key refer to https://forums.developer.nvidia.com/t/18-04-cuda-docker-image-is-broken/212892/9
|
||||
RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub 32
|
||||
RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/7fa2af80.pub
|
||||
|
||||
ARG MMENGINE="0.2.0"
|
||||
ARG MMCV="2.0.0rc1"
|
||||
ARG MMCLS="1.0.0rc2"
|
||||
|
|
@ -23,6 +27,7 @@ RUN export FORCE_CUDA=1
|
|||
|
||||
# TORCHSEVER
|
||||
RUN pip install torchserve torch-model-archiver
|
||||
RUN pip install nvgpu
|
||||
|
||||
# MMLAB
|
||||
ARG PYTORCH
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ You can switch between Chinese and English documentation in the lower-left corne
|
|||
useful_tools/verify_dataset.md
|
||||
useful_tools/log_result_analysis.md
|
||||
useful_tools/complexity_analysis.md
|
||||
useful_tools/model_serving.md
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
|
|
|||
|
|
@ -0,0 +1,87 @@
|
|||
# Torchserve Deployment
|
||||
|
||||
In order to serve an `MMClassification` model with [`TorchServe`](https://pytorch.org/serve/), you can follow the steps:
|
||||
|
||||
## 1. Convert model from MMClassification to TorchServe
|
||||
|
||||
```shell
|
||||
python tools/torchserve/mmcls2torchserve.py ${CONFIG_FILE} ${CHECKPOINT_FILE} \
|
||||
--output-folder ${MODEL_STORE} \
|
||||
--model-name ${MODEL_NAME}
|
||||
```
|
||||
|
||||
```{note}
|
||||
${MODEL_STORE} needs to be an absolute path to a folder.
|
||||
```
|
||||
|
||||
Example:
|
||||
|
||||
```shell
|
||||
python tools/torchserve/mmcls2torchserve.py \
|
||||
configs/resnet/resnet18_8xb32_in1k.py \
|
||||
checkpoints/resnet18_8xb32_in1k_20210831-fbbb1da6.pth \
|
||||
--output-folder ./checkpoints \
|
||||
--model-name resnet18_in1k
|
||||
```
|
||||
|
||||
## 2. Build `mmcls-serve` docker image
|
||||
|
||||
```shell
|
||||
docker build -t mmcls-serve:latest docker/serve/
|
||||
```
|
||||
|
||||
## 3. Run `mmcls-serve`
|
||||
|
||||
Check the official docs for [running TorchServe with docker](https://github.com/pytorch/serve/blob/master/docker/README.md#running-torchserve-in-a-production-docker-environment).
|
||||
|
||||
In order to run in GPU, you need to install [nvidia-docker](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html). You can omit the `--gpus` argument in order to run in GPU.
|
||||
|
||||
Example:
|
||||
|
||||
```shell
|
||||
docker run --rm \
|
||||
--cpus 8 \
|
||||
--gpus device=0 \
|
||||
-p8080:8080 -p8081:8081 -p8082:8082 \
|
||||
--mount type=bind,source=`realpath ./checkpoints`,target=/home/model-server/model-store \
|
||||
mmcls-serve:latest
|
||||
```
|
||||
|
||||
```{note}
|
||||
`realpath ./checkpoints` points to the absolute path of "./checkpoints", and you can replace it with the absolute path where you store torchserve models.
|
||||
```
|
||||
|
||||
[Read the docs](https://github.com/pytorch/serve/blob/master/docs/rest_api.md) about the Inference (8080), Management (8081) and Metrics (8082) APis
|
||||
|
||||
## 4. Test deployment
|
||||
|
||||
```shell
|
||||
curl http://127.0.0.1:8080/predictions/${MODEL_NAME} -T demo/demo.JPEG
|
||||
```
|
||||
|
||||
You should obtain a response similar to:
|
||||
|
||||
```json
|
||||
{
|
||||
"pred_label": 58,
|
||||
"pred_score": 0.38102269172668457,
|
||||
"pred_class": "water snake"
|
||||
}
|
||||
```
|
||||
|
||||
And you can use `test_torchserver.py` to compare result of TorchServe and PyTorch, and visualize them.
|
||||
|
||||
```shell
|
||||
python tools/torchserve/test_torchserver.py ${IMAGE_FILE} ${CONFIG_FILE} ${CHECKPOINT_FILE} ${MODEL_NAME}
|
||||
[--inference-addr ${INFERENCE_ADDR}] [--device ${DEVICE}]
|
||||
```
|
||||
|
||||
Example:
|
||||
|
||||
```shell
|
||||
python tools/torchserve/test_torchserver.py \
|
||||
demo/demo.JPEG \
|
||||
configs/resnet/resnet18_8xb32_in1k.py \
|
||||
checkpoints/resnet18_8xb32_in1k_20210831-fbbb1da6.pth \
|
||||
resnet18_in1k
|
||||
```
|
||||
|
|
@ -32,6 +32,7 @@ You can switch between Chinese and English documentation in the lower-left corne
|
|||
useful_tools/verify_dataset.md
|
||||
useful_tools/log_result_analysis.md
|
||||
useful_tools/complexity_analysis.md
|
||||
useful_tools/model_serving.md
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 1
|
||||
|
|
|
|||
|
|
@ -0,0 +1,87 @@
|
|||
# TorchServe 部署
|
||||
|
||||
为了使用 [`TorchServe`](https://pytorch.org/serve/) 部署一个 `MMClassification` 模型，需要进行以下几步：
|
||||
|
||||
## 1. 转换 MMClassification 模型至 TorchServe
|
||||
|
||||
```shell
|
||||
python tools/torchserve/mmcls2torchserve.py ${CONFIG_FILE} ${CHECKPOINT_FILE} \
|
||||
--output-folder ${MODEL_STORE} \
|
||||
--model-name ${MODEL_NAME}
|
||||
```
|
||||
|
||||
```{note}
|
||||
${MODEL_STORE} 需要是一个文件夹的绝对路径。
|
||||
```
|
||||
|
||||
示例：
|
||||
|
||||
```shell
|
||||
python tools/torchserve/mmcls2torchserve.py \
|
||||
configs/resnet/resnet18_8xb32_in1k.py \
|
||||
checkpoints/resnet18_8xb32_in1k_20210831-fbbb1da6.pth \
|
||||
--output-folder ./checkpoints \
|
||||
--model-name resnet18_in1k
|
||||
```
|
||||
|
||||
## 2. 构建 `mmcls-serve` docker 镜像
|
||||
|
||||
```shell
|
||||
docker build -t mmcls-serve:latest docker/serve/
|
||||
```
|
||||
|
||||
## 3. 运行 `mmcls-serve` 镜像
|
||||
|
||||
请参考官方文档 [基于 docker 运行 TorchServe](https://github.com/pytorch/serve/blob/master/docker/README.md#running-torchserve-in-a-production-docker-environment).
|
||||
|
||||
为了使镜像能够使用 GPU 资源，需要安装 [nvidia-docker](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html)。之后可以传递 `--gpus` 参数以在 GPU 上运。
|
||||
|
||||
示例：
|
||||
|
||||
```shell
|
||||
docker run --rm \
|
||||
--cpus 8 \
|
||||
--gpus device=0 \
|
||||
-p8080:8080 -p8081:8081 -p8082:8082 \
|
||||
--mount type=bind,source=`realpath ./checkpoints`,target=/home/model-server/model-store \
|
||||
mmcls-serve:latest
|
||||
```
|
||||
|
||||
```{note}
|
||||
`realpath ./checkpoints` 是 "./checkpoints" 的绝对路径，你可以将其替换为你保存 TorchServe 模型的目录的绝对路径。
|
||||
```
|
||||
|
||||
参考 [该文档](https://github.com/pytorch/serve/blob/master/docs/rest_api.md) 了解关于推理 (8080)，管理 (8081) 和指标 (8082) 等 API 的信息。
|
||||
|
||||
## 4. 测试部署
|
||||
|
||||
```shell
|
||||
curl http://127.0.0.1:8080/predictions/${MODEL_NAME} -T demo/demo.JPEG
|
||||
```
|
||||
|
||||
您应该获得类似于以下内容的响应：
|
||||
|
||||
```json
|
||||
{
|
||||
"pred_label": 58,
|
||||
"pred_score": 0.38102269172668457,
|
||||
"pred_class": "water snake"
|
||||
}
|
||||
```
|
||||
|
||||
另外，你也可以使用 `test_torchserver.py` 来比较 TorchServe 和 PyTorch 的结果，并进行可视化。
|
||||
|
||||
```shell
|
||||
python tools/torchserve/test_torchserver.py ${IMAGE_FILE} ${CONFIG_FILE} ${CHECKPOINT_FILE} ${MODEL_NAME}
|
||||
[--inference-addr ${INFERENCE_ADDR}] [--device ${DEVICE}]
|
||||
```
|
||||
|
||||
示例：
|
||||
|
||||
```shell
|
||||
python tools/torchserve/test_torchserver.py \
|
||||
demo/demo.JPEG \
|
||||
configs/resnet/resnet18_8xb32_in1k.py \
|
||||
checkpoints/resnet18_8xb32_in1k_20210831-fbbb1da6.pth \
|
||||
resnet18_in1k
|
||||
```
|
||||
|
|
@ -1,155 +0,0 @@
|
|||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import argparse
|
||||
import os
|
||||
import os.path as osp
|
||||
import warnings
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def get_GiB(x: int):
|
||||
"""return x GiB."""
|
||||
return x * (1 << 30)
|
||||
|
||||
|
||||
def onnx2tensorrt(onnx_file,
|
||||
trt_file,
|
||||
input_shape,
|
||||
max_batch_size,
|
||||
fp16_mode=False,
|
||||
verify=False,
|
||||
workspace_size=1):
|
||||
"""Create tensorrt engine from onnx model.
|
||||
|
||||
Args:
|
||||
onnx_file (str): Filename of the input ONNX model file.
|
||||
trt_file (str): Filename of the output TensorRT engine file.
|
||||
input_shape (list[int]): Input shape of the model.
|
||||
eg [1, 3, 224, 224].
|
||||
max_batch_size (int): Max batch size of the model.
|
||||
verify (bool, optional): Whether to verify the converted model.
|
||||
Defaults to False.
|
||||
workspace_size (int, optional): Maximum workspace of GPU.
|
||||
Defaults to 1.
|
||||
"""
|
||||
import onnx
|
||||
from mmcv.tensorrt import TRTWraper, onnx2trt, save_trt_engine
|
||||
|
||||
onnx_model = onnx.load(onnx_file)
|
||||
# create trt engine and wrapper
|
||||
assert max_batch_size >= 1
|
||||
max_shape = [max_batch_size] + list(input_shape[1:])
|
||||
opt_shape_dict = {'input': [input_shape, input_shape, max_shape]}
|
||||
max_workspace_size = get_GiB(workspace_size)
|
||||
trt_engine = onnx2trt(
|
||||
onnx_model,
|
||||
opt_shape_dict,
|
||||
fp16_mode=fp16_mode,
|
||||
max_workspace_size=max_workspace_size)
|
||||
save_dir, _ = osp.split(trt_file)
|
||||
if save_dir:
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
save_trt_engine(trt_engine, trt_file)
|
||||
print(f'Successfully created TensorRT engine: {trt_file}')
|
||||
|
||||
if verify:
|
||||
import onnxruntime as ort
|
||||
import torch
|
||||
|
||||
input_img = torch.randn(*input_shape)
|
||||
input_img_cpu = input_img.detach().cpu().numpy()
|
||||
input_img_cuda = input_img.cuda()
|
||||
|
||||
# Get results from ONNXRuntime
|
||||
session_options = ort.SessionOptions()
|
||||
sess = ort.InferenceSession(onnx_file, session_options)
|
||||
|
||||
# get input and output names
|
||||
input_names = [_.name for _ in sess.get_inputs()]
|
||||
output_names = [_.name for _ in sess.get_outputs()]
|
||||
|
||||
onnx_outputs = sess.run(None, {
|
||||
input_names[0]: input_img_cpu,
|
||||
})
|
||||
|
||||
# Get results from TensorRT
|
||||
trt_model = TRTWraper(trt_file, input_names, output_names)
|
||||
with torch.no_grad():
|
||||
trt_outputs = trt_model({input_names[0]: input_img_cuda})
|
||||
trt_outputs = [
|
||||
trt_outputs[_].detach().cpu().numpy() for _ in output_names
|
||||
]
|
||||
|
||||
# Compare results
|
||||
np.testing.assert_allclose(
|
||||
onnx_outputs[0], trt_outputs[0], rtol=1e-05, atol=1e-05)
|
||||
print('The numerical values are the same ' +
|
||||
'between ONNXRuntime and TensorRT')
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
description='Convert MMClassification models from ONNX to TensorRT')
|
||||
parser.add_argument('model', help='Filename of the input ONNX model')
|
||||
parser.add_argument(
|
||||
'--trt-file',
|
||||
type=str,
|
||||
default='tmp.trt',
|
||||
help='Filename of the output TensorRT engine')
|
||||
parser.add_argument(
|
||||
'--verify',
|
||||
action='store_true',
|
||||
help='Verify the outputs of ONNXRuntime and TensorRT')
|
||||
parser.add_argument(
|
||||
'--shape',
|
||||
type=int,
|
||||
nargs='+',
|
||||
default=[224, 224],
|
||||
help='Input size of the model')
|
||||
parser.add_argument(
|
||||
'--max-batch-size',
|
||||
type=int,
|
||||
default=1,
|
||||
help='Maximum batch size of TensorRT model.')
|
||||
parser.add_argument('--fp16', action='store_true', help='Enable fp16 mode')
|
||||
parser.add_argument(
|
||||
'--workspace-size',
|
||||
type=int,
|
||||
default=1,
|
||||
help='Max workspace size of GPU in GiB')
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
args = parse_args()
|
||||
|
||||
if len(args.shape) == 1:
|
||||
input_shape = (1, 3, args.shape[0], args.shape[0])
|
||||
elif len(args.shape) == 2:
|
||||
input_shape = (1, 3) + tuple(args.shape)
|
||||
else:
|
||||
raise ValueError('invalid input shape')
|
||||
|
||||
# Create TensorRT engine
|
||||
onnx2tensorrt(
|
||||
args.model,
|
||||
args.trt_file,
|
||||
input_shape,
|
||||
args.max_batch_size,
|
||||
fp16_mode=args.fp16,
|
||||
verify=args.verify,
|
||||
workspace_size=args.workspace_size)
|
||||
|
||||
# Following strings of text style are from colorama package
|
||||
bright_style, reset_style = '\x1b[1m', '\x1b[0m'
|
||||
red_text, blue_text = '\x1b[31m', '\x1b[34m'
|
||||
white_background = '\x1b[107m'
|
||||
|
||||
msg = white_background + bright_style + red_text
|
||||
msg += 'DeprecationWarning: This tool will be deprecated in future. '
|
||||
msg += blue_text + 'Welcome to use the unified model deployment toolbox '
|
||||
msg += 'MMDeploy: https://github.com/open-mmlab/mmdeploy'
|
||||
msg += reset_style
|
||||
warnings.warn(msg)
|
||||
|
|
@ -1,160 +0,0 @@
|
|||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import argparse
|
||||
import os
|
||||
import os.path as osp
|
||||
import warnings
|
||||
from functools import partial
|
||||
|
||||
import mmcv
|
||||
import numpy as np
|
||||
import torch
|
||||
from mmcv.runner import load_checkpoint
|
||||
from torch import nn
|
||||
|
||||
from mmcls.models import build_classifier
|
||||
|
||||
torch.manual_seed(3)
|
||||
|
||||
try:
|
||||
import coremltools as ct
|
||||
except ImportError:
|
||||
raise ImportError('Please install coremltools to enable output file.')
|
||||
|
||||
|
||||
def _demo_mm_inputs(input_shape: tuple, num_classes: int):
|
||||
"""Create a superset of inputs needed to run test or train batches.
|
||||
|
||||
Args:
|
||||
input_shape (tuple):
|
||||
input batch dimensions
|
||||
num_classes (int):
|
||||
number of semantic classes
|
||||
"""
|
||||
(N, C, H, W) = input_shape
|
||||
rng = np.random.RandomState(0)
|
||||
imgs = rng.rand(*input_shape)
|
||||
gt_labels = rng.randint(
|
||||
low=0, high=num_classes, size=(N, 1)).astype(np.uint8)
|
||||
mm_inputs = {
|
||||
'imgs': torch.FloatTensor(imgs).requires_grad_(False),
|
||||
'gt_labels': torch.LongTensor(gt_labels),
|
||||
}
|
||||
return mm_inputs
|
||||
|
||||
|
||||
def pytorch2mlmodel(model: nn.Module, input_shape: tuple, output_file: str,
|
||||
add_norm: bool, norm: dict):
|
||||
"""Export Pytorch model to mlmodel format that can be deployed in apple
|
||||
devices through torch.jit.trace and the coremltools library.
|
||||
|
||||
Optionally, embed the normalization step as a layer to the model.
|
||||
|
||||
Args:
|
||||
model (nn.Module): Pytorch model we want to export.
|
||||
input_shape (tuple): Use this input shape to construct
|
||||
the corresponding dummy input and execute the model.
|
||||
show (bool): Whether print the computation graph. Default: False.
|
||||
output_file (string): The path to where we store the output
|
||||
TorchScript model.
|
||||
add_norm (bool): Whether to embed the normalization layer to the
|
||||
output model.
|
||||
norm (dict): image normalization config for embedding it as a layer
|
||||
to the output model.
|
||||
"""
|
||||
model.cpu().eval()
|
||||
|
||||
num_classes = model.head.num_classes
|
||||
mm_inputs = _demo_mm_inputs(input_shape, num_classes)
|
||||
|
||||
imgs = mm_inputs.pop('imgs')
|
||||
img_list = [img[None, :] for img in imgs]
|
||||
model.forward = partial(model.forward, img_metas={}, return_loss=False)
|
||||
|
||||
with torch.no_grad():
|
||||
trace_model = torch.jit.trace(model, img_list[0])
|
||||
save_dir, _ = osp.split(output_file)
|
||||
if save_dir:
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
||||
if add_norm:
|
||||
means, stds = norm.mean, norm.std
|
||||
if stds.count(stds[0]) != len(stds):
|
||||
warnings.warn(f'Image std from config is {stds}. However, '
|
||||
'current version of coremltools (5.1) uses a '
|
||||
'global std rather than the channel-specific '
|
||||
'values that torchvision uses. A mean will be '
|
||||
'taken but this might tamper with the resulting '
|
||||
'model\'s predictions. For more details refer '
|
||||
'to the coreml docs on ImageType pre-processing')
|
||||
scale = np.mean(stds)
|
||||
else:
|
||||
scale = stds[0]
|
||||
|
||||
bias = [-mean / scale for mean in means]
|
||||
image_input = ct.ImageType(
|
||||
name='input_1',
|
||||
shape=input_shape,
|
||||
scale=1 / scale,
|
||||
bias=bias,
|
||||
color_layout='RGB',
|
||||
channel_first=True)
|
||||
|
||||
coreml_model = ct.convert(trace_model, inputs=[image_input])
|
||||
coreml_model.save(output_file)
|
||||
else:
|
||||
coreml_model = ct.convert(
|
||||
trace_model, inputs=[ct.TensorType(shape=input_shape)])
|
||||
coreml_model.save(output_file)
|
||||
|
||||
print(f'Successfully exported coreml model: {output_file}')
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
description='Convert MMCls to MlModel format for apple devices')
|
||||
parser.add_argument('config', help='test config file path')
|
||||
parser.add_argument('--checkpoint', help='checkpoint file', type=str)
|
||||
parser.add_argument('--output-file', type=str, default='model.mlmodel')
|
||||
parser.add_argument(
|
||||
'--shape',
|
||||
type=int,
|
||||
nargs='+',
|
||||
default=[224, 224],
|
||||
help='input image size')
|
||||
parser.add_argument(
|
||||
'--add-norm-layer',
|
||||
action='store_true',
|
||||
help='embed normalization layer to deployed model')
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
args = parse_args()
|
||||
|
||||
if len(args.shape) == 1:
|
||||
input_shape = (1, 3, args.shape[0], args.shape[0])
|
||||
elif len(args.shape) == 2:
|
||||
input_shape = (
|
||||
1,
|
||||
3,
|
||||
) + tuple(args.shape)
|
||||
else:
|
||||
raise ValueError('invalid input shape')
|
||||
|
||||
cfg = mmcv.Config.fromfile(args.config)
|
||||
cfg.model.pretrained = None
|
||||
|
||||
# build the model and load checkpoint
|
||||
classifier = build_classifier(cfg.model)
|
||||
|
||||
if args.checkpoint:
|
||||
load_checkpoint(classifier, args.checkpoint, map_location='cpu')
|
||||
|
||||
# convert model to mlmodel file
|
||||
pytorch2mlmodel(
|
||||
classifier,
|
||||
input_shape,
|
||||
output_file=args.output_file,
|
||||
add_norm=args.add_norm_layer,
|
||||
norm=cfg.img_norm_cfg)
|
||||
|
|
@ -1,246 +0,0 @@
|
|||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import argparse
|
||||
import warnings
|
||||
from functools import partial
|
||||
|
||||
import mmcv
|
||||
import numpy as np
|
||||
import onnxruntime as rt
|
||||
import torch
|
||||
from mmcv.onnx import register_extra_symbolics
|
||||
from mmcv.runner import load_checkpoint
|
||||
|
||||
from mmcls.models import build_classifier
|
||||
|
||||
torch.manual_seed(3)
|
||||
|
||||
|
||||
def _demo_mm_inputs(input_shape, num_classes):
|
||||
"""Create a superset of inputs needed to run test or train batches.
|
||||
|
||||
Args:
|
||||
input_shape (tuple):
|
||||
input batch dimensions
|
||||
num_classes (int):
|
||||
number of semantic classes
|
||||
"""
|
||||
(N, C, H, W) = input_shape
|
||||
rng = np.random.RandomState(0)
|
||||
imgs = rng.rand(*input_shape)
|
||||
gt_labels = rng.randint(
|
||||
low=0, high=num_classes, size=(N, 1)).astype(np.uint8)
|
||||
mm_inputs = {
|
||||
'imgs': torch.FloatTensor(imgs).requires_grad_(True),
|
||||
'gt_labels': torch.LongTensor(gt_labels),
|
||||
}
|
||||
return mm_inputs
|
||||
|
||||
|
||||
def pytorch2onnx(model,
|
||||
input_shape,
|
||||
opset_version=11,
|
||||
dynamic_export=False,
|
||||
show=False,
|
||||
output_file='tmp.onnx',
|
||||
do_simplify=False,
|
||||
verify=False):
|
||||
"""Export Pytorch model to ONNX model and verify the outputs are same
|
||||
between Pytorch and ONNX.
|
||||
|
||||
Args:
|
||||
model (nn.Module): Pytorch model we want to export.
|
||||
input_shape (tuple): Use this input shape to construct
|
||||
the corresponding dummy input and execute the model.
|
||||
opset_version (int): The onnx op version. Default: 11.
|
||||
show (bool): Whether print the computation graph. Default: False.
|
||||
output_file (string): The path to where we store the output ONNX model.
|
||||
Default: `tmp.onnx`.
|
||||
verify (bool): Whether compare the outputs between Pytorch and ONNX.
|
||||
Default: False.
|
||||
"""
|
||||
model.cpu().eval()
|
||||
|
||||
if hasattr(model.head, 'num_classes'):
|
||||
num_classes = model.head.num_classes
|
||||
# Some backbones use `num_classes=-1` to disable top classifier.
|
||||
elif getattr(model.backbone, 'num_classes', -1) > 0:
|
||||
num_classes = model.backbone.num_classes
|
||||
else:
|
||||
raise AttributeError('Cannot find "num_classes" in both head and '
|
||||
'backbone, please check the config file.')
|
||||
|
||||
mm_inputs = _demo_mm_inputs(input_shape, num_classes)
|
||||
|
||||
imgs = mm_inputs.pop('imgs')
|
||||
img_list = [img[None, :] for img in imgs]
|
||||
|
||||
# replace original forward function
|
||||
origin_forward = model.forward
|
||||
model.forward = partial(model.forward, img_metas={}, return_loss=False)
|
||||
register_extra_symbolics(opset_version)
|
||||
|
||||
# support dynamic shape export
|
||||
if dynamic_export:
|
||||
dynamic_axes = {
|
||||
'input': {
|
||||
0: 'batch',
|
||||
2: 'width',
|
||||
3: 'height'
|
||||
},
|
||||
'probs': {
|
||||
0: 'batch'
|
||||
}
|
||||
}
|
||||
else:
|
||||
dynamic_axes = {}
|
||||
|
||||
with torch.no_grad():
|
||||
torch.onnx.export(
|
||||
model, (img_list, ),
|
||||
output_file,
|
||||
input_names=['input'],
|
||||
output_names=['probs'],
|
||||
export_params=True,
|
||||
keep_initializers_as_inputs=True,
|
||||
dynamic_axes=dynamic_axes,
|
||||
verbose=show,
|
||||
opset_version=opset_version)
|
||||
print(f'Successfully exported ONNX model: {output_file}')
|
||||
model.forward = origin_forward
|
||||
|
||||
if do_simplify:
|
||||
import onnx
|
||||
import onnxsim
|
||||
from mmcv import digit_version
|
||||
|
||||
min_required_version = '0.3.0'
|
||||
assert digit_version(mmcv.__version__) >= digit_version(
|
||||
min_required_version
|
||||
), f'Requires to install onnx-simplify>={min_required_version}'
|
||||
|
||||
if dynamic_axes:
|
||||
input_shape = (input_shape[0], input_shape[1], input_shape[2] * 2,
|
||||
input_shape[3] * 2)
|
||||
else:
|
||||
input_shape = (input_shape[0], input_shape[1], input_shape[2],
|
||||
input_shape[3])
|
||||
imgs = _demo_mm_inputs(input_shape, model.head.num_classes).pop('imgs')
|
||||
input_dic = {'input': imgs.detach().cpu().numpy()}
|
||||
input_shape_dic = {'input': list(input_shape)}
|
||||
|
||||
model_opt, check_ok = onnxsim.simplify(
|
||||
output_file,
|
||||
input_shapes=input_shape_dic,
|
||||
input_data=input_dic,
|
||||
dynamic_input_shape=dynamic_export)
|
||||
if check_ok:
|
||||
onnx.save(model_opt, output_file)
|
||||
print(f'Successfully simplified ONNX model: {output_file}')
|
||||
else:
|
||||
print('Failed to simplify ONNX model.')
|
||||
if verify:
|
||||
# check by onnx
|
||||
import onnx
|
||||
onnx_model = onnx.load(output_file)
|
||||
onnx.checker.check_model(onnx_model)
|
||||
|
||||
# test the dynamic model
|
||||
if dynamic_export:
|
||||
dynamic_test_inputs = _demo_mm_inputs(
|
||||
(input_shape[0], input_shape[1], input_shape[2] * 2,
|
||||
input_shape[3] * 2), model.head.num_classes)
|
||||
imgs = dynamic_test_inputs.pop('imgs')
|
||||
img_list = [img[None, :] for img in imgs]
|
||||
|
||||
# check the numerical value
|
||||
# get pytorch output
|
||||
pytorch_result = model(img_list, img_metas={}, return_loss=False)[0]
|
||||
|
||||
# get onnx output
|
||||
input_all = [node.name for node in onnx_model.graph.input]
|
||||
input_initializer = [
|
||||
node.name for node in onnx_model.graph.initializer
|
||||
]
|
||||
net_feed_input = list(set(input_all) - set(input_initializer))
|
||||
assert (len(net_feed_input) == 1)
|
||||
sess = rt.InferenceSession(output_file)
|
||||
onnx_result = sess.run(
|
||||
None, {net_feed_input[0]: img_list[0].detach().numpy()})[0]
|
||||
if not np.allclose(pytorch_result, onnx_result):
|
||||
raise ValueError(
|
||||
'The outputs are different between Pytorch and ONNX')
|
||||
print('The outputs are same between Pytorch and ONNX')
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description='Convert MMCls to ONNX')
|
||||
parser.add_argument('config', help='test config file path')
|
||||
parser.add_argument('--checkpoint', help='checkpoint file', default=None)
|
||||
parser.add_argument('--show', action='store_true', help='show onnx graph')
|
||||
parser.add_argument(
|
||||
'--verify', action='store_true', help='verify the onnx model')
|
||||
parser.add_argument('--output-file', type=str, default='tmp.onnx')
|
||||
parser.add_argument('--opset-version', type=int, default=11)
|
||||
parser.add_argument(
|
||||
'--simplify',
|
||||
action='store_true',
|
||||
help='Whether to simplify onnx model.')
|
||||
parser.add_argument(
|
||||
'--shape',
|
||||
type=int,
|
||||
nargs='+',
|
||||
default=[224, 224],
|
||||
help='input image size')
|
||||
parser.add_argument(
|
||||
'--dynamic-export',
|
||||
action='store_true',
|
||||
help='Whether to export ONNX with dynamic input shape. \
|
||||
Defaults to False.')
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
args = parse_args()
|
||||
|
||||
if len(args.shape) == 1:
|
||||
input_shape = (1, 3, args.shape[0], args.shape[0])
|
||||
elif len(args.shape) == 2:
|
||||
input_shape = (
|
||||
1,
|
||||
3,
|
||||
) + tuple(args.shape)
|
||||
else:
|
||||
raise ValueError('invalid input shape')
|
||||
|
||||
cfg = mmcv.Config.fromfile(args.config)
|
||||
cfg.model.pretrained = None
|
||||
|
||||
# build the model and load checkpoint
|
||||
classifier = build_classifier(cfg.model)
|
||||
|
||||
if args.checkpoint:
|
||||
load_checkpoint(classifier, args.checkpoint, map_location='cpu')
|
||||
|
||||
# convert model to onnx file
|
||||
pytorch2onnx(
|
||||
classifier,
|
||||
input_shape,
|
||||
opset_version=args.opset_version,
|
||||
show=args.show,
|
||||
dynamic_export=args.dynamic_export,
|
||||
output_file=args.output_file,
|
||||
do_simplify=args.simplify,
|
||||
verify=args.verify)
|
||||
|
||||
# Following strings of text style are from colorama package
|
||||
bright_style, reset_style = '\x1b[1m', '\x1b[0m'
|
||||
red_text, blue_text = '\x1b[31m', '\x1b[34m'
|
||||
white_background = '\x1b[107m'
|
||||
|
||||
msg = white_background + bright_style + red_text
|
||||
msg += 'DeprecationWarning: This tool will be deprecated in future. '
|
||||
msg += blue_text + 'Welcome to use the unified model deployment toolbox '
|
||||
msg += 'MMDeploy: https://github.com/open-mmlab/mmdeploy'
|
||||
msg += reset_style
|
||||
warnings.warn(msg)
|
||||
|
|
@ -1,139 +0,0 @@
|
|||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import argparse
|
||||
import os
|
||||
import os.path as osp
|
||||
from functools import partial
|
||||
|
||||
import mmcv
|
||||
import numpy as np
|
||||
import torch
|
||||
from mmcv.runner import load_checkpoint
|
||||
from torch import nn
|
||||
|
||||
from mmcls.models import build_classifier
|
||||
|
||||
torch.manual_seed(3)
|
||||
|
||||
|
||||
def _demo_mm_inputs(input_shape: tuple, num_classes: int):
|
||||
"""Create a superset of inputs needed to run test or train batches.
|
||||
|
||||
Args:
|
||||
input_shape (tuple):
|
||||
input batch dimensions
|
||||
num_classes (int):
|
||||
number of semantic classes
|
||||
"""
|
||||
(N, C, H, W) = input_shape
|
||||
rng = np.random.RandomState(0)
|
||||
imgs = rng.rand(*input_shape)
|
||||
gt_labels = rng.randint(
|
||||
low=0, high=num_classes, size=(N, 1)).astype(np.uint8)
|
||||
mm_inputs = {
|
||||
'imgs': torch.FloatTensor(imgs).requires_grad_(False),
|
||||
'gt_labels': torch.LongTensor(gt_labels),
|
||||
}
|
||||
return mm_inputs
|
||||
|
||||
|
||||
def pytorch2torchscript(model: nn.Module, input_shape: tuple, output_file: str,
|
||||
verify: bool):
|
||||
"""Export Pytorch model to TorchScript model through torch.jit.trace and
|
||||
verify the outputs are same between Pytorch and TorchScript.
|
||||
|
||||
Args:
|
||||
model (nn.Module): Pytorch model we want to export.
|
||||
input_shape (tuple): Use this input shape to construct
|
||||
the corresponding dummy input and execute the model.
|
||||
show (bool): Whether print the computation graph. Default: False.
|
||||
output_file (string): The path to where we store the output
|
||||
TorchScript model.
|
||||
verify (bool): Whether compare the outputs between Pytorch
|
||||
and TorchScript through loading generated output_file.
|
||||
"""
|
||||
model.cpu().eval()
|
||||
|
||||
num_classes = model.head.num_classes
|
||||
mm_inputs = _demo_mm_inputs(input_shape, num_classes)
|
||||
|
||||
imgs = mm_inputs.pop('imgs')
|
||||
img_list = [img[None, :] for img in imgs]
|
||||
|
||||
# replace original forward function
|
||||
origin_forward = model.forward
|
||||
model.forward = partial(model.forward, img_metas={}, return_loss=False)
|
||||
|
||||
with torch.no_grad():
|
||||
trace_model = torch.jit.trace(model, img_list[0])
|
||||
save_dir, _ = osp.split(output_file)
|
||||
if save_dir:
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
trace_model.save(output_file)
|
||||
print(f'Successfully exported TorchScript model: {output_file}')
|
||||
model.forward = origin_forward
|
||||
|
||||
if verify:
|
||||
# load by torch.jit
|
||||
jit_model = torch.jit.load(output_file)
|
||||
|
||||
# check the numerical value
|
||||
# get pytorch output
|
||||
pytorch_result = model(img_list, img_metas={}, return_loss=False)[0]
|
||||
|
||||
# get jit output
|
||||
jit_result = jit_model(img_list[0])[0].detach().numpy()
|
||||
if not np.allclose(pytorch_result, jit_result):
|
||||
raise ValueError(
|
||||
'The outputs are different between Pytorch and TorchScript')
|
||||
print('The outputs are same between Pytorch and TorchScript')
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
description='Convert MMCls to TorchScript')
|
||||
parser.add_argument('config', help='test config file path')
|
||||
parser.add_argument('--checkpoint', help='checkpoint file', type=str)
|
||||
parser.add_argument(
|
||||
'--verify',
|
||||
action='store_true',
|
||||
help='verify the TorchScript model',
|
||||
default=False)
|
||||
parser.add_argument('--output-file', type=str, default='tmp.pt')
|
||||
parser.add_argument(
|
||||
'--shape',
|
||||
type=int,
|
||||
nargs='+',
|
||||
default=[224, 224],
|
||||
help='input image size')
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
args = parse_args()
|
||||
|
||||
if len(args.shape) == 1:
|
||||
input_shape = (1, 3, args.shape[0], args.shape[0])
|
||||
elif len(args.shape) == 2:
|
||||
input_shape = (
|
||||
1,
|
||||
3,
|
||||
) + tuple(args.shape)
|
||||
else:
|
||||
raise ValueError('invalid input shape')
|
||||
|
||||
cfg = mmcv.Config.fromfile(args.config)
|
||||
cfg.model.pretrained = None
|
||||
|
||||
# build the model and load checkpoint
|
||||
classifier = build_classifier(cfg.model)
|
||||
|
||||
if args.checkpoint:
|
||||
load_checkpoint(classifier, args.checkpoint, map_location='cpu')
|
||||
|
||||
# convert model to TorchScript file
|
||||
pytorch2torchscript(
|
||||
classifier,
|
||||
input_shape,
|
||||
output_file=args.output_file,
|
||||
verify=args.verify)
|
||||
|
|
@ -1,128 +0,0 @@
|
|||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import argparse
|
||||
import warnings
|
||||
|
||||
import mmcv
|
||||
import numpy as np
|
||||
from mmcv import DictAction
|
||||
from mmcv.parallel import MMDataParallel
|
||||
|
||||
from mmcls.apis import single_gpu_test
|
||||
from mmcls.datasets import build_dataloader, build_dataset
|
||||
from mmcls.engine.export import ONNXRuntimeClassifier, TensorRTClassifier
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
description='Test (and eval) an ONNX model using ONNXRuntime.')
|
||||
parser.add_argument('config', help='model config file')
|
||||
parser.add_argument('model', help='filename of the input ONNX model')
|
||||
parser.add_argument(
|
||||
'--backend',
|
||||
help='Backend of the model.',
|
||||
choices=['onnxruntime', 'tensorrt'])
|
||||
parser.add_argument(
|
||||
'--out', type=str, help='output result file in pickle format')
|
||||
parser.add_argument(
|
||||
'--cfg-options',
|
||||
nargs='+',
|
||||
action=DictAction,
|
||||
help='override some settings in the used config, the key-value pair '
|
||||
'in xxx=yyy format will be merged into config file.')
|
||||
parser.add_argument(
|
||||
'--metrics',
|
||||
type=str,
|
||||
nargs='+',
|
||||
help='evaluation metrics, which depends on the dataset, e.g., '
|
||||
'"accuracy", "precision", "recall", "f1_score", "support" for single '
|
||||
'label dataset, and "mAP", "CP", "CR", "CF1", "OP", "OR", "OF1" for '
|
||||
'multi-label dataset')
|
||||
parser.add_argument(
|
||||
'--metric-options',
|
||||
nargs='+',
|
||||
action=DictAction,
|
||||
default={},
|
||||
help='custom options for evaluation, the key-value pair in xxx=yyy '
|
||||
'format will be parsed as a dict metric_options for dataset.evaluate()'
|
||||
' function.')
|
||||
parser.add_argument('--show', action='store_true', help='show results')
|
||||
parser.add_argument(
|
||||
'--show-dir', help='directory where painted images will be saved')
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
if args.out is not None and not args.out.endswith(('.pkl', '.pickle')):
|
||||
raise ValueError('The output file must be a pkl file.')
|
||||
|
||||
cfg = mmcv.Config.fromfile(args.config)
|
||||
if args.cfg_options is not None:
|
||||
cfg.merge_from_dict(args.cfg_options)
|
||||
|
||||
# build dataset and dataloader
|
||||
dataset = build_dataset(cfg.data.test)
|
||||
data_loader = build_dataloader(
|
||||
dataset,
|
||||
samples_per_gpu=cfg.data.samples_per_gpu,
|
||||
workers_per_gpu=cfg.data.workers_per_gpu,
|
||||
shuffle=False,
|
||||
round_up=False)
|
||||
|
||||
# build onnxruntime model and run inference.
|
||||
if args.backend == 'onnxruntime':
|
||||
model = ONNXRuntimeClassifier(
|
||||
args.model, class_names=dataset.CLASSES, device_id=0)
|
||||
elif args.backend == 'tensorrt':
|
||||
model = TensorRTClassifier(
|
||||
args.model, class_names=dataset.CLASSES, device_id=0)
|
||||
else:
|
||||
print('Unknown backend: {}.'.format(args.model))
|
||||
exit(1)
|
||||
|
||||
model = MMDataParallel(model, device_ids=[0])
|
||||
model.CLASSES = dataset.CLASSES
|
||||
outputs = single_gpu_test(model, data_loader, args.show, args.show_dir)
|
||||
|
||||
if args.metrics:
|
||||
results = dataset.evaluate(outputs, args.metrics, args.metric_options)
|
||||
for k, v in results.items():
|
||||
print(f'\n{k} : {v:.2f}')
|
||||
else:
|
||||
warnings.warn('Evaluation metrics are not specified.')
|
||||
scores = np.vstack(outputs)
|
||||
pred_score = np.max(scores, axis=1)
|
||||
pred_label = np.argmax(scores, axis=1)
|
||||
pred_class = [dataset.CLASSES[lb] for lb in pred_label]
|
||||
results = {
|
||||
'pred_score': pred_score,
|
||||
'pred_label': pred_label,
|
||||
'pred_class': pred_class
|
||||
}
|
||||
if not args.out:
|
||||
print('\nthe predicted result for the first element is '
|
||||
f'pred_score = {pred_score[0]:.2f}, '
|
||||
f'pred_label = {pred_label[0]} '
|
||||
f'and pred_class = {pred_class[0]}. '
|
||||
'Specify --out to save all results to files.')
|
||||
if args.out:
|
||||
print(f'\nwriting results to {args.out}')
|
||||
mmcv.dump(results, args.out)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
|
||||
# Following strings of text style are from colorama package
|
||||
bright_style, reset_style = '\x1b[1m', '\x1b[0m'
|
||||
red_text, blue_text = '\x1b[31m', '\x1b[34m'
|
||||
white_background = '\x1b[107m'
|
||||
|
||||
msg = white_background + bright_style + red_text
|
||||
msg += 'DeprecationWarning: This tool will be deprecated in future. '
|
||||
msg += blue_text + 'Welcome to use the unified model deployment toolbox '
|
||||
msg += 'MMDeploy: https://github.com/open-mmlab/mmdeploy'
|
||||
msg += reset_style
|
||||
warnings.warn(msg)
|
||||
|
|
@ -3,14 +3,15 @@ from argparse import ArgumentParser, Namespace
|
|||
from pathlib import Path
|
||||
from tempfile import TemporaryDirectory
|
||||
|
||||
from mmengine.config import Config
|
||||
from mmengine.utils import mkdir_or_exist
|
||||
import mmengine
|
||||
|
||||
try:
|
||||
from model_archiver.model_packaging import package_model
|
||||
from model_archiver.model_packaging_utils import ModelExportUtils
|
||||
except ImportError:
|
||||
package_model = None
|
||||
raise ImportError(
|
||||
'Please run `pip install torchserve torch-model-archiver"` to '
|
||||
'install required third-party libraries.')
|
||||
|
||||
|
||||
def mmcls2torchserve(
|
||||
|
|
@ -44,9 +45,9 @@ def mmcls2torchserve(
|
|||
If True, if there is an existing `{model_name}.mar`
|
||||
file under `output_folder` it will be overwritten.
|
||||
"""
|
||||
mkdir_or_exist(output_folder)
|
||||
mmengine.mkdir_or_exist(output_folder)
|
||||
|
||||
config = Config.fromfile(config_file)
|
||||
config = mmengine.Config.fromfile(config_file)
|
||||
|
||||
with TemporaryDirectory() as tmpdir:
|
||||
config.dump(f'{tmpdir}/config.py')
|
||||
|
|
@ -4,7 +4,7 @@ from argparse import ArgumentParser
|
|||
import numpy as np
|
||||
import requests
|
||||
|
||||
from mmcls.apis import inference_model, init_model, show_result_pyplot
|
||||
from mmcls.apis import inference_model, init_model
|
||||
|
||||
|
||||
def parse_args():
|
||||
|
|
@ -27,14 +27,12 @@ def main(args):
|
|||
# Inference single image by native apis.
|
||||
model = init_model(args.config, args.checkpoint, device=args.device)
|
||||
model_result = inference_model(model, args.img)
|
||||
show_result_pyplot(model, args.img, model_result, title='pytorch_result')
|
||||
|
||||
# Inference single image by torchserve engine.
|
||||
url = 'http://' + args.inference_addr + '/predictions/' + args.model_name
|
||||
with open(args.img, 'rb') as image:
|
||||
response = requests.post(url, image)
|
||||
server_result = response.json()
|
||||
show_result_pyplot(model, args.img, server_result, title='server_result')
|
||||
|
||||
assert np.allclose(model_result['pred_score'], server_result['pred_score'])
|
||||
print('Test complete, the results of PyTorch and TorchServe are the same.')
|
||||
Loading…
Reference in New Issue