Model Demo: ConvNeXt Segmentation

NOTE: The Jupyter Notebook below is included in the Chimera SDK and can be run interactively by running the following CLI command:

$ quadric sdk notebook

From the Jupyter Notebook window in your browser, select the notebook named /quadric/sdk-cli/examples/models/convnext/segmenter/convnext_segmention.ipynb.

ConvNeXt Segmentation

Abstract

A vanilla Vision Transformers (ViTs) faces difficulties when applied to general computer vision tasks such as object detection and semantic segmentation. It is the hierarchical Transformers (e.g., Swin Transformers) that reintroduced several ConvNet priors, making Transformers practically viable as a generic vision backbone and demonstrating remarkable performance on a wide variety of vision tasks. However, the effectiveness of such hybrid approaches is still largely credited to the intrinsic superiority of Transformers, rather than the inherent inductive biases of convolutions.

By gradually "moderniz"ing a standard ResNet toward the design of a vision Transformer, several key components are discovered that contribute to the performance difference along the way. The outcome of this exploration is a family of pure ConvNet models dubbed ConvNeXt. Constructed entirely from standard ConvNet modules, ConvNeXts compete favorably with Transformers in terms of accuracy and scalability, achieving 87.8% ImageNet top-1 accuracy and outperforming Swin Transformers on COCO detection and ADE20K segmentation, while maintaining the simplicity and efficiency of standard ConvNets.

Further information can be found in the following Paper.

• Zhuang Liu, Hanzi Mao, Chao-Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, Saining Xie. A ConvNet for the 2020s

In this nodebook, we are going to experiment ConvNeXt for ADE20K segmentation.

Set-Up

Install and Import Packages

!pip3 install -r ../../../requirements.txt -q
!mim install -r ../../../requirements_mim.txt -q

import gc
import numpy as np
import onnx
from onnxruntime import InferenceSession
import onnxsim
import os
from pathlib import Path
from PIL import Image
import random
import shutil
import torch
from torch.utils.data import Subset
from torchvision.transforms import Compose, Normalize, ToTensor
import urllib
import zipfile

from examples.models.zoo.zoo_utils import (
    download_file,
    get_mmdeploy_parameters,
    get_transforms_and_dataset,
    install_openmmlab_github,
    onnx_check_and_simplify,
    split_onnx,
)
import examples.models.zoo.fix_registry
from examples.models.zoo.segmentors_zoo.segmentors_utils import (
    get_mmseg_parameters,
)
from tvm.contrib.epu.chimera_job.chimera_job import ChimeraJob
from sdk_cli.lib.inference import InferenceEngine, batch_inference
from sdk_cli.lib.quantize import (
    QuantizedONNXModel,
    quantize_onnx_model,
)
from sdk_cli.utils.datasets.ADE20K import ADE20K150CLASSES
from sdk_cli.utils.models.mmdeploy import MMDeployVariant
from sdk_cli.utils.models.mmseg import MMSegModelVariant
from sdk_cli.utils.transforms import ResizePad
from sdk_cli.visualizers.layouter import (
    DetectorsLayouter,
    VisualizeHandle,
    VisualizeObject,
    get_visualize_handle,
)
from tvm.contrib.epu.onnx_util import cut_onnx


CUR_DIR = os.path.abspath(os.getcwd())

Git Clone MMSegmentation and MMDeploy

github_names = [
    ("mmsegmentation", "1.2.2", None),
    ("mmdeploy", "1.3.1", None),
]
symlink_list = [
    ("mmseg", "mmsegmentation"),
    ("configs", "mmsegmentation"),
    ("model-index.yml", "mmsegmentation"),
]
install_openmmlab_github(github_names, symlink_list)

import mmseg

os.chdir(Path(CUR_DIR) / "mmdeploy")

from mmdeploy.apis import torch2onnx
import torch

os.chdir(CUR_DIR)

Create Configs

config_text = """_base_ = ['./segmentation_static.py', '../_base_/backends/ascend.py']

onnx_config = dict(input_shape=[%s, %s])
backend_config = dict(
    model_inputs=[dict(input_shapes=dict(input=[1, 3, %s, %s]))])
"""
config_path = "./mmdeploy/configs/mmseg/segmentation_onnxruntime_static-%sx%s.py"

for model_input_size in [(640, 640), (1600, 928)]:
    if not Path(config_path % model_input_size[::-1]).exists():
        with Path(config_path % model_input_size[::-1]).open("w") as fd:
            fd.writelines(config_text % (*model_input_size, *model_input_size[::-1]))
            print(f"{config_path % model_input_size[::-1]} is stored")

Define Model

In this nodebook, the following models can be experiment. Here we are going to experiment convnext_tiny_upernet_640x640 of which input size is (640, 640).

• convnext_tiny_upernet_640x640
• convnext_tiny_upernet_928x1600
• convnext_small_upernet_640x640
• convnext_small_upernet_928x1600

## Select target model.
MODEL_NAME = MMSegModelVariant.convnext_tiny_upernet_640x640

mmseg_parameters = get_mmseg_parameters(MODEL_NAME)
model_input_size = mmseg_parameters.model_input_size
mmdeploy_variant = MMDeployVariant.mmsegmentation

mmdeploy_parameters = get_mmdeploy_parameters(mmdeploy_variant, mmseg_parameters)

Load Pretrained Model

weights = Path(mmdeploy_parameters.weights_url).name

if not Path(weights).exists():
    download_file(mmdeploy_parameters.weights_url, weights)
    print(f"{weights} is downloaded.")

Export ONNX

onnx_file = f"{MODEL_NAME}.onnx"

## Convert the model to ONNX
torch2onnx(
    "../../../common/calibration/face/daniel_maverick.png",
    work_dir=".",
    save_file=onnx_file,
    deploy_cfg=mmdeploy_parameters.deploy_config,
    model_cfg=mmdeploy_parameters.config,
    model_checkpoint=weights,
    device="cpu",
)

del torch, torch2onnx, mmseg
gc.collect();

Simplify the onnx file.

onnx_model = onnx_check_and_simplify(onnx.load(onnx_file), mmseg_parameters.skip_shape_inference)
onnx.save(onnx_model, onnx_file)

print(f"ONNX is exported and simplified as {onnx_file}")

Extract Backbone and Head

The target network consists of two networks.

• ConvNeXt as a backbone
• UPerNet as a head

We are going to extract ConvNeXt as the backbone and UPerNet as the head.

session = InferenceSession(onnx.load(onnx_file).SerializeToString())
input_edges = [inp.name for inp in session.get_inputs()]
output_edges = [oup.name for oup in session.get_outputs()]

target_onnx_file, head_onnx_file = split_onnx(
    onnx_file,
    mmseg_parameters,
    input_edges=input_edges,
    output_edges=output_edges,
)

gc.collect();

Quantization

transforms, coco_like_subset_of_dataset = get_transforms_and_dataset(mmdeploy_parameters)

quantized_onnx_model: QuantizedONNXModel = quantize_onnx_model(
    target_onnx_file,
    coco_like_subset_of_dataset,
    asymmetric_activation=True,
)
print(
    f"quantized onnx: {quantized_onnx_model.model_path}, tranges file: {quantized_onnx_model.tensor_ranges_path}"
)

del coco_like_subset_of_dataset
gc.collect();

Compilation

## Compilation.
cgc_job = ChimeraJob(
    model_p=str(quantized_onnx_model.model_path),
    trange_file=str(quantized_onnx_model.tensor_ranges_path),
)

cgc_job.compile(quiet=True)
print(cgc_job)

Demo

Inference

all_image_paths = [
    "../../../common/calibration/coco-like/17790319373_bd19b24cfc_k.jpeg",
    "../../../common/calibration/coco-like/19064748793_bb942deea1_k.jpeg",
    "../../../common/calibration/coco-like/24274813513_0cfd2ce6d0_k.jpeg",
]

all_images = []
for image_path in all_image_paths:
    image = Image.open(image_path)
    transformed_image = transforms(image)
    all_images.append(np.expand_dims(transformed_image.numpy(), axis=0).astype(np.float32))

engines = {
    InferenceEngine.CHIMERA_ORT_INT8: cgc_job,
    InferenceEngine.CHIMERA_ISS_INT8: cgc_job,
}

head_session = InferenceSession(onnx.load(head_onnx_file).SerializeToString())

outputs_per_inference_engine = {}
THREADS = min(len(all_images), 6)
for inference_engine, engine in engines.items():
    all_outputs = batch_inference(
        inference_engine,
        engine,
        all_images,
        head_session=head_session,
        threads=THREADS,
    )
    outputs_per_inference_engine[inference_engine] = all_outputs

Run Statistics

print(cgc_job)
cgc_job.plot_run_statistics()

Display Inference Results

%matplotlib inline

visualize_object = VisualizeObject.SemanticSegmentation
visualize_handle = get_visualize_handle(visualize_object, mmdeploy_parameters, random_seed=17)

for i in range(len(all_images)):
    layouter = DetectorsLayouter(
        visualize_handle,
        image=all_image_paths[i],
        classes=ADE20K150CLASSES,
    )
    for inference_engine, all_outputs in outputs_per_inference_engine.items():
        layouter.add_data(
            bboxes=None,
            masks=all_outputs[i][0][0],
            title=f"{MODEL_NAME}: {str(inference_engine).upper()}",
        )

    layouter.display()

Citation

@article{liu2022convnet,
  title={A ConvNet for the 2020s},
  author={Liu, Zhuang and Mao, Hanzi and Wu, Chao-Yuan and Feichtenhofer, Christoph and Darrell, Trevor and Xie, Saining},
  journal={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
  year={2022}
}