# Copyright (c) Facebook, Inc. and its affiliates. import io import numpy as np import os import re import tempfile import unittest from typing import Callable import torch import torch.onnx.symbolic_helper as sym_help from packaging import version from torch._C import ListType from torch.onnx import register_custom_op_symbolic from detectron2 import model_zoo from detectron2.config import CfgNode, LazyConfig, instantiate from detectron2.data import DatasetCatalog from detectron2.data.detection_utils import read_image from detectron2.modeling import build_model from detectron2.structures import Boxes, Instances, ROIMasks from detectron2.utils.file_io import PathManager """ Internal utilities for tests. Don't use except for writing tests. """ def get_model_no_weights(config_path): """ Like model_zoo.get, but do not load any weights (even pretrained) """ cfg = model_zoo.get_config(config_path) if isinstance(cfg, CfgNode): if not torch.cuda.is_available(): cfg.MODEL.DEVICE = "cpu" return build_model(cfg) else: return instantiate(cfg.model) def random_boxes(num_boxes, max_coord=100, device="cpu"): """ Create a random Nx4 boxes tensor, with coordinates < max_coord. """ boxes = torch.rand(num_boxes, 4, device=device) * (max_coord * 0.5) boxes.clamp_(min=1.0) # tiny boxes cause numerical instability in box regression # Note: the implementation of this function in torchvision is: # boxes[:, 2:] += torch.rand(N, 2) * 100 # but it does not guarantee non-negative widths/heights constraints: # boxes[:, 2] >= boxes[:, 0] and boxes[:, 3] >= boxes[:, 1]: boxes[:, 2:] += boxes[:, :2] return boxes def get_sample_coco_image(tensor=True): """ Args: tensor (bool): if True, returns 3xHxW tensor. else, returns a HxWx3 numpy array. Returns: an image, in BGR color. """ try: file_name = DatasetCatalog.get("coco_2017_val_100")[0]["file_name"] if not PathManager.exists(file_name): raise FileNotFoundError() except IOError: # for public CI to run file_name = PathManager.get_local_path( "http://images.cocodataset.org/train2017/000000000009.jpg" ) ret = read_image(file_name, format="BGR") if tensor: ret = torch.from_numpy(np.ascontiguousarray(ret.transpose(2, 0, 1))) return ret def convert_scripted_instances(instances): """ Convert a scripted Instances object to a regular :class:`Instances` object """ assert hasattr( instances, "image_size" ), f"Expect an Instances object, but got {type(instances)}!" ret = Instances(instances.image_size) for name in instances._field_names: val = getattr(instances, "_" + name, None) if val is not None: ret.set(name, val) return ret def assert_instances_allclose(input, other, *, rtol=1e-5, msg="", size_as_tensor=False): """ Args: input, other (Instances): size_as_tensor: compare image_size of the Instances as tensors (instead of tuples). Useful for comparing outputs of tracing. """ if not isinstance(input, Instances): input = convert_scripted_instances(input) if not isinstance(other, Instances): other = convert_scripted_instances(other) if not msg: msg = "Two Instances are different! " else: msg = msg.rstrip() + " " size_error_msg = msg + f"image_size is {input.image_size} vs. {other.image_size}!" if size_as_tensor: assert torch.equal( torch.tensor(input.image_size), torch.tensor(other.image_size) ), size_error_msg else: assert input.image_size == other.image_size, size_error_msg fields = sorted(input.get_fields().keys()) fields_other = sorted(other.get_fields().keys()) assert fields == fields_other, msg + f"Fields are {fields} vs {fields_other}!" for f in fields: val1, val2 = input.get(f), other.get(f) if isinstance(val1, (Boxes, ROIMasks)): # boxes in the range of O(100) and can have a larger tolerance assert torch.allclose(val1.tensor, val2.tensor, atol=100 * rtol), ( msg + f"Field {f} differs too much!" ) elif isinstance(val1, torch.Tensor): if val1.dtype.is_floating_point: mag = torch.abs(val1).max().cpu().item() assert torch.allclose(val1, val2, atol=mag * rtol), ( msg + f"Field {f} differs too much!" ) else: assert torch.equal(val1, val2), msg + f"Field {f} is different!" else: raise ValueError(f"Don't know how to compare type {type(val1)}") def reload_script_model(module): """ Save a jit module and load it back. Similar to the `getExportImportCopy` function in torch/testing/ """ buffer = io.BytesIO() torch.jit.save(module, buffer) buffer.seek(0) return torch.jit.load(buffer) def reload_lazy_config(cfg): """ Save an object by LazyConfig.save and load it back. This is used to test that a config still works the same after serialization/deserialization. """ with tempfile.TemporaryDirectory(prefix="detectron2") as d: fname = os.path.join(d, "d2_cfg_test.yaml") LazyConfig.save(cfg, fname) return LazyConfig.load(fname) def min_torch_version(min_version: str) -> bool: """ Returns True when torch's version is at least `min_version`. """ try: import torch except ImportError: return False installed_version = version.parse(torch.__version__.split("+")[0]) min_version = version.parse(min_version) return installed_version >= min_version def has_dynamic_axes(onnx_model): """ Return True when all ONNX input/output have only dynamic axes for all ranks """ return all( not dim.dim_param.isnumeric() for inp in onnx_model.graph.input for dim in inp.type.tensor_type.shape.dim ) and all( not dim.dim_param.isnumeric() for out in onnx_model.graph.output for dim in out.type.tensor_type.shape.dim ) def register_custom_op_onnx_export( opname: str, symbolic_fn: Callable, opset_version: int, min_version: str ) -> None: """ Register `symbolic_fn` as PyTorch's symbolic `opname`-`opset_version` for ONNX export. The registration is performed only when current PyTorch's version is < `min_version.` IMPORTANT: symbolic must be manually unregistered after the caller function returns """ if min_torch_version(min_version): return register_custom_op_symbolic(opname, symbolic_fn, opset_version) print(f"_register_custom_op_onnx_export({opname}, {opset_version}) succeeded.") def unregister_custom_op_onnx_export(opname: str, opset_version: int, min_version: str) -> None: """ Unregister PyTorch's symbolic `opname`-`opset_version` for ONNX export. The un-registration is performed only when PyTorch's version is < `min_version` IMPORTANT: The symbolic must have been manually registered by the caller, otherwise the incorrect symbolic may be unregistered instead. """ # TODO: _unregister_custom_op_symbolic is introduced PyTorch>=1.10 # Remove after PyTorch 1.10+ is used by ALL detectron2's CI try: from torch.onnx import unregister_custom_op_symbolic as _unregister_custom_op_symbolic except ImportError: def _unregister_custom_op_symbolic(symbolic_name, opset_version): import torch.onnx.symbolic_registry as sym_registry from torch.onnx.symbolic_helper import _onnx_main_opset, _onnx_stable_opsets def _get_ns_op_name_from_custom_op(symbolic_name): try: from torch.onnx.utils import get_ns_op_name_from_custom_op ns, op_name = get_ns_op_name_from_custom_op(symbolic_name) except ImportError as import_error: if not bool( re.match(r"^[a-zA-Z0-9-_]*::[a-zA-Z-_]+[a-zA-Z0-9-_]*$", symbolic_name) ): raise ValueError( f"Invalid symbolic name {symbolic_name}. Must be `domain::name`" ) from import_error ns, op_name = symbolic_name.split("::") if ns == "onnx": raise ValueError(f"{ns} domain cannot be modified.") from import_error if ns == "aten": ns = "" return ns, op_name def _unregister_op(opname: str, domain: str, version: int): try: sym_registry.unregister_op(op_name, ns, ver) except AttributeError as attribute_error: if sym_registry.is_registered_op(opname, domain, version): del sym_registry._registry[(domain, version)][opname] if not sym_registry._registry[(domain, version)]: del sym_registry._registry[(domain, version)] else: raise RuntimeError( f"The opname {opname} is not registered." ) from attribute_error ns, op_name = _get_ns_op_name_from_custom_op(symbolic_name) for ver in _onnx_stable_opsets + [_onnx_main_opset]: if ver >= opset_version: _unregister_op(op_name, ns, ver) if min_torch_version(min_version): return _unregister_custom_op_symbolic(opname, opset_version) print(f"_unregister_custom_op_onnx_export({opname}, {opset_version}) succeeded.") skipIfOnCPUCI = unittest.skipIf( os.environ.get("CI") and not torch.cuda.is_available(), "The test is too slow on CPUs and will be executed on CircleCI's GPU jobs.", ) def skipIfUnsupportedMinOpsetVersion(min_opset_version, current_opset_version=None): """ Skips tests for ONNX Opset versions older than min_opset_version. """ def skip_dec(func): def wrapper(self): try: opset_version = self.opset_version except AttributeError: opset_version = current_opset_version if opset_version < min_opset_version: raise unittest.SkipTest( f"Unsupported opset_version {opset_version}" f", required is {min_opset_version}" ) return func(self) return wrapper return skip_dec def skipIfUnsupportedMinTorchVersion(min_version): """ Skips tests for PyTorch versions older than min_version. """ reason = f"module 'torch' has __version__ {torch.__version__}" f", required is: {min_version}" return unittest.skipIf(not min_torch_version(min_version), reason) # TODO: Remove after PyTorch 1.11.1+ is used by detectron2's CI def _pytorch1111_symbolic_opset9_to(g, self, *args): """aten::to() symbolic that must be used for testing with PyTorch < 1.11.1.""" def is_aten_to_device_only(args): if len(args) == 4: # aten::to(Tensor, Device, bool, bool, memory_format) return ( args[0].node().kind() == "prim::device" or args[0].type().isSubtypeOf(ListType.ofInts()) or ( sym_help._is_value(args[0]) and args[0].node().kind() == "onnx::Constant" and isinstance(args[0].node()["value"], str) ) ) elif len(args) == 5: # aten::to(Tensor, Device, ScalarType, bool, bool, memory_format) # When dtype is None, this is a aten::to(device) call dtype = sym_help._get_const(args[1], "i", "dtype") return dtype is None elif len(args) in (6, 7): # aten::to(Tensor, ScalarType, Layout, Device, bool, bool, memory_format) # aten::to(Tensor, ScalarType, Layout, Device, bool, bool, bool, memory_format) # When dtype is None, this is a aten::to(device) call dtype = sym_help._get_const(args[0], "i", "dtype") return dtype is None return False # ONNX doesn't have a concept of a device, so we ignore device-only casts if is_aten_to_device_only(args): return self if len(args) == 4: # TestONNXRuntime::test_ones_bool shows args[0] of aten::to can be onnx::Constant[Tensor] # In this case, the constant value is a tensor not int, # so sym_help._maybe_get_const(args[0], 'i') would not work. dtype = args[0] if sym_help._is_value(args[0]) and args[0].node().kind() == "onnx::Constant": tval = args[0].node()["value"] if isinstance(tval, torch.Tensor): if len(tval.shape) == 0: tval = tval.item() dtype = int(tval) else: dtype = tval if sym_help._is_value(dtype) or isinstance(dtype, torch.Tensor): # aten::to(Tensor, Tensor, bool, bool, memory_format) dtype = args[0].type().scalarType() return g.op("Cast", self, to_i=sym_help.cast_pytorch_to_onnx[dtype]) else: # aten::to(Tensor, ScalarType, bool, bool, memory_format) # memory_format is ignored return g.op("Cast", self, to_i=sym_help.scalar_type_to_onnx[dtype]) elif len(args) == 5: # aten::to(Tensor, Device, ScalarType, bool, bool, memory_format) dtype = sym_help._get_const(args[1], "i", "dtype") # memory_format is ignored return g.op("Cast", self, to_i=sym_help.scalar_type_to_onnx[dtype]) elif len(args) == 6: # aten::to(Tensor, ScalarType, Layout, Device, bool, bool, memory_format) dtype = sym_help._get_const(args[0], "i", "dtype") # Layout, device and memory_format are ignored return g.op("Cast", self, to_i=sym_help.scalar_type_to_onnx[dtype]) elif len(args) == 7: # aten::to(Tensor, ScalarType, Layout, Device, bool, bool, bool, memory_format) dtype = sym_help._get_const(args[0], "i", "dtype") # Layout, device and memory_format are ignored return g.op("Cast", self, to_i=sym_help.scalar_type_to_onnx[dtype]) else: return sym_help._onnx_unsupported("Unknown aten::to signature") # TODO: Remove after PyTorch 1.11.1+ is used by detectron2's CI def _pytorch1111_symbolic_opset9_repeat_interleave(g, self, repeats, dim=None, output_size=None): # from torch.onnx.symbolic_helper import ScalarType from torch.onnx.symbolic_opset9 import expand, unsqueeze input = self # if dim is None flatten # By default, use the flattened input array, and return a flat output array if sym_help._is_none(dim): input = sym_help._reshape_helper(g, self, g.op("Constant", value_t=torch.tensor([-1]))) dim = 0 else: dim = sym_help._maybe_get_scalar(dim) repeats_dim = sym_help._get_tensor_rank(repeats) repeats_sizes = sym_help._get_tensor_sizes(repeats) input_sizes = sym_help._get_tensor_sizes(input) if repeats_dim is None: raise RuntimeError( "Unsupported: ONNX export of repeat_interleave for unknown " "repeats rank." ) if repeats_sizes is None: raise RuntimeError( "Unsupported: ONNX export of repeat_interleave for unknown " "repeats size." ) if input_sizes is None: raise RuntimeError( "Unsupported: ONNX export of repeat_interleave for unknown " "input size." ) input_sizes_temp = input_sizes.copy() for idx, input_size in enumerate(input_sizes): if input_size is None: input_sizes[idx], input_sizes_temp[idx] = 0, -1 # Cases where repeats is an int or single value tensor if repeats_dim == 0 or (repeats_dim == 1 and repeats_sizes[0] == 1): if not sym_help._is_tensor(repeats): repeats = g.op("Constant", value_t=torch.LongTensor(repeats)) if input_sizes[dim] == 0: return sym_help._onnx_opset_unsupported_detailed( "repeat_interleave", 9, 13, "Unsupported along dimension with unknown input size", ) else: reps = input_sizes[dim] repeats = expand(g, repeats, g.op("Constant", value_t=torch.tensor([reps])), None) # Cases where repeats is a 1 dim Tensor elif repeats_dim == 1: if input_sizes[dim] == 0: return sym_help._onnx_opset_unsupported_detailed( "repeat_interleave", 9, 13, "Unsupported along dimension with unknown input size", ) if repeats_sizes[0] is None: return sym_help._onnx_opset_unsupported_detailed( "repeat_interleave", 9, 13, "Unsupported for cases with dynamic repeats" ) assert ( repeats_sizes[0] == input_sizes[dim] ), "repeats must have the same size as input along dim" reps = repeats_sizes[0] else: raise RuntimeError("repeats must be 0-dim or 1-dim tensor") final_splits = list() r_splits = sym_help._repeat_interleave_split_helper(g, repeats, reps, 0) if isinstance(r_splits, torch._C.Value): r_splits = [r_splits] i_splits = sym_help._repeat_interleave_split_helper(g, input, reps, dim) if isinstance(i_splits, torch._C.Value): i_splits = [i_splits] input_sizes[dim], input_sizes_temp[dim] = -1, 1 for idx, r_split in enumerate(r_splits): i_split = unsqueeze(g, i_splits[idx], dim + 1) r_concat = [ g.op("Constant", value_t=torch.LongTensor(input_sizes_temp[: dim + 1])), r_split, g.op("Constant", value_t=torch.LongTensor(input_sizes_temp[dim + 1 :])), ] r_concat = g.op("Concat", *r_concat, axis_i=0) i_split = expand(g, i_split, r_concat, None) i_split = sym_help._reshape_helper( g, i_split, g.op("Constant", value_t=torch.LongTensor(input_sizes)), allowzero=0, ) final_splits.append(i_split) return g.op("Concat", *final_splits, axis_i=dim)