# Copyright (c) Facebook, Inc. and its affiliates.
import io
import unittest
import warnings
import torch
from torch.hub import _check_module_exists
from detectron2 import model_zoo
from detectron2.config import get_cfg
from detectron2.export import STABLE_ONNX_OPSET_VERSION
from detectron2.export.flatten import TracingAdapter
from detectron2.export.torchscript_patch import patch_builtin_len
from detectron2.layers import ShapeSpec
from detectron2.modeling import build_model
from detectron2.modeling.roi_heads import KRCNNConvDeconvUpsampleHead
from detectron2.structures import Boxes, Instances
from detectron2.utils.testing import (
_pytorch1111_symbolic_opset9_repeat_interleave,
_pytorch1111_symbolic_opset9_to,
get_sample_coco_image,
has_dynamic_axes,
random_boxes,
register_custom_op_onnx_export,
skipIfOnCPUCI,
skipIfUnsupportedMinOpsetVersion,
skipIfUnsupportedMinTorchVersion,
unregister_custom_op_onnx_export,
)
@unittest.skipIf(not _check_module_exists("onnx"), "ONNX not installed.")
@skipIfUnsupportedMinTorchVersion("1.10")
class TestONNXTracingExport(unittest.TestCase):
opset_version = STABLE_ONNX_OPSET_VERSION
def testMaskRCNNFPN(self):
def inference_func(model, images):
with warnings.catch_warnings(record=True):
inputs = [{"image": image} for image in images]
inst = model.inference(inputs, do_postprocess=False)[0]
return [{"instances": inst}]
self._test_model_zoo_from_config_path(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml", inference_func
)
@skipIfOnCPUCI
def testMaskRCNNC4(self):
def inference_func(model, image):
inputs = [{"image": image}]
return model.inference(inputs, do_postprocess=False)[0]
self._test_model_zoo_from_config_path(
"COCO-InstanceSegmentation/mask_rcnn_R_50_C4_3x.yaml", inference_func
)
@skipIfOnCPUCI
def testCascadeRCNN(self):
def inference_func(model, image):
inputs = [{"image": image}]
return model.inference(inputs, do_postprocess=False)[0]
self._test_model_zoo_from_config_path(
"Misc/cascade_mask_rcnn_R_50_FPN_3x.yaml", inference_func
)
def testRetinaNet(self):
def inference_func(model, image):
return model.forward([{"image": image}])[0]["instances"]
self._test_model_zoo_from_config_path(
"COCO-Detection/retinanet_R_50_FPN_3x.yaml", inference_func
)
@skipIfOnCPUCI
def testMaskRCNNFPN_batched(self):
def inference_func(model, image1, image2):
inputs = [{"image": image1}, {"image": image2}]
return model.inference(inputs, do_postprocess=False)
self._test_model_zoo_from_config_path(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml", inference_func, batch=2
)
@skipIfUnsupportedMinOpsetVersion(16, STABLE_ONNX_OPSET_VERSION)
@skipIfUnsupportedMinTorchVersion("1.11.1")
def testMaskRCNNFPN_with_postproc(self):
def inference_func(model, image):
inputs = [{"image": image, "height": image.shape[1], "width": image.shape[2]}]
return model.inference(inputs, do_postprocess=True)[0]["instances"]
self._test_model_zoo_from_config_path(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml",
inference_func,
)
def testKeypointHead(self):
class M(torch.nn.Module):
def __init__(self):
super().__init__()
self.model = KRCNNConvDeconvUpsampleHead(
ShapeSpec(channels=4, height=14, width=14), num_keypoints=17, conv_dims=(4,)
)
def forward(self, x, predbox1, predbox2):
inst = [
Instances((100, 100), pred_boxes=Boxes(predbox1)),
Instances((100, 100), pred_boxes=Boxes(predbox2)),
]
ret = self.model(x, inst)
return tuple(x.pred_keypoints for x in ret)
model = M()
model.eval()
def gen_input(num1, num2):
feat = torch.randn((num1 + num2, 4, 14, 14))
box1 = random_boxes(num1)
box2 = random_boxes(num2)
return feat, box1, box2
with patch_builtin_len():
onnx_model = self._test_model(
model,
gen_input(1, 2),
input_names=["features", "pred_boxes", "pred_classes"],
output_names=["box1", "box2"],
dynamic_axes={
"features": {0: "batch", 1: "static_four", 2: "height", 3: "width"},
"pred_boxes": {0: "batch", 1: "static_four"},
"pred_classes": {0: "batch", 1: "static_four"},
"box1": {0: "num_instance", 1: "K", 2: "static_three"},
"box2": {0: "num_instance", 1: "K", 2: "static_three"},
},
)
# Although ONNX models are not executable by PyTorch to verify
# support of batches with different sizes, we can verify model's IR
# does not hard-code input and/or output shapes.
# TODO: Add tests with different batch sizes when detectron2's CI
# support ONNX Runtime backend.
assert has_dynamic_axes(onnx_model)
################################################################################
# Testcase internals - DO NOT add tests below this point
################################################################################
def setUp(self):
register_custom_op_onnx_export("::to", _pytorch1111_symbolic_opset9_to, 9, "1.11.1")
register_custom_op_onnx_export(
"::repeat_interleave",
_pytorch1111_symbolic_opset9_repeat_interleave,
9,
"1.11.1",
)
def tearDown(self):
unregister_custom_op_onnx_export("::to", 9, "1.11.1")
unregister_custom_op_onnx_export("::repeat_interleave", 9, "1.11.1")
def _test_model(
self,
model,
inputs,
inference_func=None,
opset_version=STABLE_ONNX_OPSET_VERSION,
save_onnx_graph_path=None,
**export_kwargs,
):
# Not imported in the beginning of file to prevent runtime errors
# for environments without ONNX.
# This testcase checks dependencies before running
import onnx # isort:skip
f = io.BytesIO()
adapter_model = TracingAdapter(model, inputs, inference_func)
adapter_model.eval()
with torch.no_grad():
try:
torch.onnx.enable_log()
except AttributeError:
# Older ONNX versions does not have this API
pass
torch.onnx.export(
adapter_model,
adapter_model.flattened_inputs,
f,
training=torch.onnx.TrainingMode.EVAL,
opset_version=opset_version,
verbose=True,
**export_kwargs,
)
onnx_model = onnx.load_from_string(f.getvalue())
assert onnx_model is not None
if save_onnx_graph_path:
onnx.save(onnx_model, save_onnx_graph_path)
return onnx_model
def _test_model_zoo_from_config_path(
self,
config_path,
inference_func,
batch=1,
opset_version=STABLE_ONNX_OPSET_VERSION,
save_onnx_graph_path=None,
**export_kwargs,
):
model = model_zoo.get(config_path, trained=True)
image = get_sample_coco_image()
inputs = tuple(image.clone() for _ in range(batch))
return self._test_model(
model, inputs, inference_func, opset_version, save_onnx_graph_path, **export_kwargs
)
def _test_model_from_config_path(
self,
config_path,
inference_func,
batch=1,
opset_version=STABLE_ONNX_OPSET_VERSION,
save_onnx_graph_path=None,
**export_kwargs,
):
from projects.PointRend import point_rend # isort:skip
cfg = get_cfg()
cfg.DATALOADER.NUM_WORKERS = 0
point_rend.add_pointrend_config(cfg)
cfg.merge_from_file(config_path)
cfg.freeze()
model = build_model(cfg)
image = get_sample_coco_image()
inputs = tuple(image.clone() for _ in range(batch))
return self._test_model(
model, inputs, inference_func, opset_version, save_onnx_graph_path, **export_kwargs
)