Arulkumar03/Wheat_HEAD_Detection_Counting_ComputerVision_Model
0
1# Copyright (c) Facebook, Inc. and its affiliates.2import io3import numpy as np4import os5import re6import tempfile7import unittest8from typing import Callable9import torch10import torch.onnx.symbolic_helper as sym_help11from packaging import version12from torch._C import ListType13from torch.onnx import register_custom_op_symbolic14 15from detectron2 import model_zoo16from detectron2.config import CfgNode, LazyConfig, instantiate17from detectron2.data import DatasetCatalog18from detectron2.data.detection_utils import read_image19from detectron2.modeling import build_model20from detectron2.structures import Boxes, Instances, ROIMasks21from detectron2.utils.file_io import PathManager22 23 24"""25Internal utilities for tests. Don't use except for writing tests.26"""27 28 29def get_model_no_weights(config_path):30 """31 Like model_zoo.get, but do not load any weights (even pretrained)32 """33 cfg = model_zoo.get_config(config_path)34 if isinstance(cfg, CfgNode):35 if not torch.cuda.is_available():36 cfg.MODEL.DEVICE = "cpu"37 return build_model(cfg)38 else:39 return instantiate(cfg.model)40 41 42def random_boxes(num_boxes, max_coord=100, device="cpu"):43 """44 Create a random Nx4 boxes tensor, with coordinates < max_coord.45 """46 boxes = torch.rand(num_boxes, 4, device=device) * (max_coord * 0.5)47 boxes.clamp_(min=1.0) # tiny boxes cause numerical instability in box regression48 # Note: the implementation of this function in torchvision is:49 # boxes[:, 2:] += torch.rand(N, 2) * 10050 # but it does not guarantee non-negative widths/heights constraints:51 # boxes[:, 2] >= boxes[:, 0] and boxes[:, 3] >= boxes[:, 1]:52 boxes[:, 2:] += boxes[:, :2]53 return boxes54 55 56def get_sample_coco_image(tensor=True):57 """58 Args:59 tensor (bool): if True, returns 3xHxW tensor.60 else, returns a HxWx3 numpy array.61 62 Returns:63 an image, in BGR color.64 """65 try:66 file_name = DatasetCatalog.get("coco_2017_val_100")[0]["file_name"]67 if not PathManager.exists(file_name):68 raise FileNotFoundError()69 except IOError:70 # for public CI to run71 file_name = PathManager.get_local_path(72 "http://images.cocodataset.org/train2017/000000000009.jpg"73 )74 ret = read_image(file_name, format="BGR")75 if tensor:76 ret = torch.from_numpy(np.ascontiguousarray(ret.transpose(2, 0, 1)))77 return ret78 79 80def convert_scripted_instances(instances):81 """82 Convert a scripted Instances object to a regular :class:`Instances` object83 """84 assert hasattr(85 instances, "image_size"86 ), f"Expect an Instances object, but got {type(instances)}!"87 ret = Instances(instances.image_size)88 for name in instances._field_names:89 val = getattr(instances, "_" + name, None)90 if val is not None:91 ret.set(name, val)92 return ret93 94 95def assert_instances_allclose(input, other, *, rtol=1e-5, msg="", size_as_tensor=False):96 """97 Args:98 input, other (Instances):99 size_as_tensor: compare image_size of the Instances as tensors (instead of tuples).100 Useful for comparing outputs of tracing.101 """102 if not isinstance(input, Instances):103 input = convert_scripted_instances(input)104 if not isinstance(other, Instances):105 other = convert_scripted_instances(other)106 107 if not msg:108 msg = "Two Instances are different! "109 else:110 msg = msg.rstrip() + " "111 112 size_error_msg = msg + f"image_size is {input.image_size} vs. {other.image_size}!"113 if size_as_tensor:114 assert torch.equal(115 torch.tensor(input.image_size), torch.tensor(other.image_size)116 ), size_error_msg117 else:118 assert input.image_size == other.image_size, size_error_msg119 fields = sorted(input.get_fields().keys())120 fields_other = sorted(other.get_fields().keys())121 assert fields == fields_other, msg + f"Fields are {fields} vs {fields_other}!"122 123 for f in fields:124 val1, val2 = input.get(f), other.get(f)125 if isinstance(val1, (Boxes, ROIMasks)):126 # boxes in the range of O(100) and can have a larger tolerance127 assert torch.allclose(val1.tensor, val2.tensor, atol=100 * rtol), (128 msg + f"Field {f} differs too much!"129 )130 elif isinstance(val1, torch.Tensor):131 if val1.dtype.is_floating_point:132 mag = torch.abs(val1).max().cpu().item()133 assert torch.allclose(val1, val2, atol=mag * rtol), (134 msg + f"Field {f} differs too much!"135 )136 else:137 assert torch.equal(val1, val2), msg + f"Field {f} is different!"138 else:139 raise ValueError(f"Don't know how to compare type {type(val1)}")140 141 142def reload_script_model(module):143 """144 Save a jit module and load it back.145 Similar to the `getExportImportCopy` function in torch/testing/146 """147 buffer = io.BytesIO()148 torch.jit.save(module, buffer)149 buffer.seek(0)150 return torch.jit.load(buffer)151 152 153def reload_lazy_config(cfg):154 """155 Save an object by LazyConfig.save and load it back.156 This is used to test that a config still works the same after157 serialization/deserialization.158 """159 with tempfile.TemporaryDirectory(prefix="detectron2") as d:160 fname = os.path.join(d, "d2_cfg_test.yaml")161 LazyConfig.save(cfg, fname)162 return LazyConfig.load(fname)163 164 165def min_torch_version(min_version: str) -> bool:166 """167 Returns True when torch's version is at least `min_version`.168 """169 try:170 import torch171 except ImportError:172 return False173 174 installed_version = version.parse(torch.__version__.split("+")[0])175 min_version = version.parse(min_version)176 return installed_version >= min_version177 178 179def has_dynamic_axes(onnx_model):180 """181 Return True when all ONNX input/output have only dynamic axes for all ranks182 """183 return all(184 not dim.dim_param.isnumeric()185 for inp in onnx_model.graph.input186 for dim in inp.type.tensor_type.shape.dim187 ) and all(188 not dim.dim_param.isnumeric()189 for out in onnx_model.graph.output190 for dim in out.type.tensor_type.shape.dim191 )192 193 194def register_custom_op_onnx_export(195 opname: str, symbolic_fn: Callable, opset_version: int, min_version: str196) -> None:197 """198 Register `symbolic_fn` as PyTorch's symbolic `opname`-`opset_version` for ONNX export.199 The registration is performed only when current PyTorch's version is < `min_version.`200 IMPORTANT: symbolic must be manually unregistered after the caller function returns201 """202 if min_torch_version(min_version):203 return204 register_custom_op_symbolic(opname, symbolic_fn, opset_version)205 print(f"_register_custom_op_onnx_export({opname}, {opset_version}) succeeded.")206 207 208def unregister_custom_op_onnx_export(opname: str, opset_version: int, min_version: str) -> None:209 """210 Unregister PyTorch's symbolic `opname`-`opset_version` for ONNX export.211 The un-registration is performed only when PyTorch's version is < `min_version`212 IMPORTANT: The symbolic must have been manually registered by the caller, otherwise213 the incorrect symbolic may be unregistered instead.214 """215 216 # TODO: _unregister_custom_op_symbolic is introduced PyTorch>=1.10217 # Remove after PyTorch 1.10+ is used by ALL detectron2's CI218 try:219 from torch.onnx import unregister_custom_op_symbolic as _unregister_custom_op_symbolic220 except ImportError:221 222 def _unregister_custom_op_symbolic(symbolic_name, opset_version):223 import torch.onnx.symbolic_registry as sym_registry224 from torch.onnx.symbolic_helper import _onnx_main_opset, _onnx_stable_opsets225 226 def _get_ns_op_name_from_custom_op(symbolic_name):227 try:228 from torch.onnx.utils import get_ns_op_name_from_custom_op229 230 ns, op_name = get_ns_op_name_from_custom_op(symbolic_name)231 except ImportError as import_error:232 if not bool(233 re.match(r"^[a-zA-Z0-9-_]*::[a-zA-Z-_]+[a-zA-Z0-9-_]*$", symbolic_name)234 ):235 raise ValueError(236 f"Invalid symbolic name {symbolic_name}. Must be `domain::name`"237 ) from import_error238 239 ns, op_name = symbolic_name.split("::")240 if ns == "onnx":241 raise ValueError(f"{ns} domain cannot be modified.") from import_error242 243 if ns == "aten":244 ns = ""245 246 return ns, op_name247 248 def _unregister_op(opname: str, domain: str, version: int):249 try:250 sym_registry.unregister_op(op_name, ns, ver)251 except AttributeError as attribute_error:252 if sym_registry.is_registered_op(opname, domain, version):253 del sym_registry._registry[(domain, version)][opname]254 if not sym_registry._registry[(domain, version)]:255 del sym_registry._registry[(domain, version)]256 else:257 raise RuntimeError(258 f"The opname {opname} is not registered."259 ) from attribute_error260 261 ns, op_name = _get_ns_op_name_from_custom_op(symbolic_name)262 for ver in _onnx_stable_opsets + [_onnx_main_opset]:263 if ver >= opset_version:264 _unregister_op(op_name, ns, ver)265 266 if min_torch_version(min_version):267 return268 _unregister_custom_op_symbolic(opname, opset_version)269 print(f"_unregister_custom_op_onnx_export({opname}, {opset_version}) succeeded.")270 271 272skipIfOnCPUCI = unittest.skipIf(273 os.environ.get("CI") and not torch.cuda.is_available(),274 "The test is too slow on CPUs and will be executed on CircleCI's GPU jobs.",275)276 277 278def skipIfUnsupportedMinOpsetVersion(min_opset_version, current_opset_version=None):279 """280 Skips tests for ONNX Opset versions older than min_opset_version.281 """282 283 def skip_dec(func):284 def wrapper(self):285 try:286 opset_version = self.opset_version287 except AttributeError:288 opset_version = current_opset_version289 if opset_version < min_opset_version:290 raise unittest.SkipTest(291 f"Unsupported opset_version {opset_version}"292 f", required is {min_opset_version}"293 )294 return func(self)295 296 return wrapper297 298 return skip_dec299 300 301def skipIfUnsupportedMinTorchVersion(min_version):302 """303 Skips tests for PyTorch versions older than min_version.304 """305 reason = f"module 'torch' has __version__ {torch.__version__}" f", required is: {min_version}"306 return unittest.skipIf(not min_torch_version(min_version), reason)307 308 309# TODO: Remove after PyTorch 1.11.1+ is used by detectron2's CI310def _pytorch1111_symbolic_opset9_to(g, self, *args):311 """aten::to() symbolic that must be used for testing with PyTorch < 1.11.1."""312 313 def is_aten_to_device_only(args):314 if len(args) == 4:315 # aten::to(Tensor, Device, bool, bool, memory_format)316 return (317 args[0].node().kind() == "prim::device"318 or args[0].type().isSubtypeOf(ListType.ofInts())319 or (320 sym_help._is_value(args[0])321 and args[0].node().kind() == "onnx::Constant"322 and isinstance(args[0].node()["value"], str)323 )324 )325 elif len(args) == 5:326 # aten::to(Tensor, Device, ScalarType, bool, bool, memory_format)327 # When dtype is None, this is a aten::to(device) call328 dtype = sym_help._get_const(args[1], "i", "dtype")329 return dtype is None330 elif len(args) in (6, 7):331 # aten::to(Tensor, ScalarType, Layout, Device, bool, bool, memory_format)332 # aten::to(Tensor, ScalarType, Layout, Device, bool, bool, bool, memory_format)333 # When dtype is None, this is a aten::to(device) call334 dtype = sym_help._get_const(args[0], "i", "dtype")335 return dtype is None336 return False337 338 # ONNX doesn't have a concept of a device, so we ignore device-only casts339 if is_aten_to_device_only(args):340 return self341 342 if len(args) == 4:343 # TestONNXRuntime::test_ones_bool shows args[0] of aten::to can be onnx::Constant[Tensor]344 # In this case, the constant value is a tensor not int,345 # so sym_help._maybe_get_const(args[0], 'i') would not work.346 dtype = args[0]347 if sym_help._is_value(args[0]) and args[0].node().kind() == "onnx::Constant":348 tval = args[0].node()["value"]349 if isinstance(tval, torch.Tensor):350 if len(tval.shape) == 0:351 tval = tval.item()352 dtype = int(tval)353 else:354 dtype = tval355 356 if sym_help._is_value(dtype) or isinstance(dtype, torch.Tensor):357 # aten::to(Tensor, Tensor, bool, bool, memory_format)358 dtype = args[0].type().scalarType()359 return g.op("Cast", self, to_i=sym_help.cast_pytorch_to_onnx[dtype])360 else:361 # aten::to(Tensor, ScalarType, bool, bool, memory_format)362 # memory_format is ignored363 return g.op("Cast", self, to_i=sym_help.scalar_type_to_onnx[dtype])364 elif len(args) == 5:365 # aten::to(Tensor, Device, ScalarType, bool, bool, memory_format)366 dtype = sym_help._get_const(args[1], "i", "dtype")367 # memory_format is ignored368 return g.op("Cast", self, to_i=sym_help.scalar_type_to_onnx[dtype])369 elif len(args) == 6:370 # aten::to(Tensor, ScalarType, Layout, Device, bool, bool, memory_format)371 dtype = sym_help._get_const(args[0], "i", "dtype")372 # Layout, device and memory_format are ignored373 return g.op("Cast", self, to_i=sym_help.scalar_type_to_onnx[dtype])374 elif len(args) == 7:375 # aten::to(Tensor, ScalarType, Layout, Device, bool, bool, bool, memory_format)376 dtype = sym_help._get_const(args[0], "i", "dtype")377 # Layout, device and memory_format are ignored378 return g.op("Cast", self, to_i=sym_help.scalar_type_to_onnx[dtype])379 else:380 return sym_help._onnx_unsupported("Unknown aten::to signature")381 382 383# TODO: Remove after PyTorch 1.11.1+ is used by detectron2's CI384def _pytorch1111_symbolic_opset9_repeat_interleave(g, self, repeats, dim=None, output_size=None):385 386 # from torch.onnx.symbolic_helper import ScalarType387 from torch.onnx.symbolic_opset9 import expand, unsqueeze388 389 input = self390 # if dim is None flatten391 # By default, use the flattened input array, and return a flat output array392 if sym_help._is_none(dim):393 input = sym_help._reshape_helper(g, self, g.op("Constant", value_t=torch.tensor([-1])))394 dim = 0395 else:396 dim = sym_help._maybe_get_scalar(dim)397 398 repeats_dim = sym_help._get_tensor_rank(repeats)399 repeats_sizes = sym_help._get_tensor_sizes(repeats)400 input_sizes = sym_help._get_tensor_sizes(input)401 if repeats_dim is None:402 raise RuntimeError(403 "Unsupported: ONNX export of repeat_interleave for unknown " "repeats rank."404 )405 if repeats_sizes is None:406 raise RuntimeError(407 "Unsupported: ONNX export of repeat_interleave for unknown " "repeats size."408 )409 if input_sizes is None:410 raise RuntimeError(411 "Unsupported: ONNX export of repeat_interleave for unknown " "input size."412 )413 414 input_sizes_temp = input_sizes.copy()415 for idx, input_size in enumerate(input_sizes):416 if input_size is None:417 input_sizes[idx], input_sizes_temp[idx] = 0, -1418 419 # Cases where repeats is an int or single value tensor420 if repeats_dim == 0 or (repeats_dim == 1 and repeats_sizes[0] == 1):421 if not sym_help._is_tensor(repeats):422 repeats = g.op("Constant", value_t=torch.LongTensor(repeats))423 if input_sizes[dim] == 0:424 return sym_help._onnx_opset_unsupported_detailed(425 "repeat_interleave",426 9,427 13,428 "Unsupported along dimension with unknown input size",429 )430 else:431 reps = input_sizes[dim]432 repeats = expand(g, repeats, g.op("Constant", value_t=torch.tensor([reps])), None)433 434 # Cases where repeats is a 1 dim Tensor435 elif repeats_dim == 1:436 if input_sizes[dim] == 0:437 return sym_help._onnx_opset_unsupported_detailed(438 "repeat_interleave",439 9,440 13,441 "Unsupported along dimension with unknown input size",442 )443 if repeats_sizes[0] is None:444 return sym_help._onnx_opset_unsupported_detailed(445 "repeat_interleave", 9, 13, "Unsupported for cases with dynamic repeats"446 )447 assert (448 repeats_sizes[0] == input_sizes[dim]449 ), "repeats must have the same size as input along dim"450 reps = repeats_sizes[0]451 else:452 raise RuntimeError("repeats must be 0-dim or 1-dim tensor")453 454 final_splits = list()455 r_splits = sym_help._repeat_interleave_split_helper(g, repeats, reps, 0)456 if isinstance(r_splits, torch._C.Value):457 r_splits = [r_splits]458 i_splits = sym_help._repeat_interleave_split_helper(g, input, reps, dim)459 if isinstance(i_splits, torch._C.Value):460 i_splits = [i_splits]461 input_sizes[dim], input_sizes_temp[dim] = -1, 1462 for idx, r_split in enumerate(r_splits):463 i_split = unsqueeze(g, i_splits[idx], dim + 1)464 r_concat = [465 g.op("Constant", value_t=torch.LongTensor(input_sizes_temp[: dim + 1])),466 r_split,467 g.op("Constant", value_t=torch.LongTensor(input_sizes_temp[dim + 1 :])),468 ]469 r_concat = g.op("Concat", *r_concat, axis_i=0)470 i_split = expand(g, i_split, r_concat, None)471 i_split = sym_help._reshape_helper(472 g,473 i_split,474 g.op("Constant", value_t=torch.LongTensor(input_sizes)),475 allowzero=0,476 )477 final_splits.append(i_split)478 return g.op("Concat", *final_splits, axis_i=dim)479