Building with TVM relay

[Tomakko]

Hi @zhiqwang ,

i am still working on deploying the yolov5 rt model with TVM. It works out fine, when i use the VM compile

  from tvm.runtime.vm import VirtualMachine
  with tvm.transform.PassContext(opt_level=3, disabled_pass=["FoldScaleAxis"]):
      vm_exec = relay.vm.compile(mod, target=build_target, params=params)

However i need to deploy it as a lib.so. Therefore i want to compile as

build_target = 'llvm'
with tvm.transform.PassContext(opt_level=3, disabled_pass=["FoldScaleAxis"]):
    graph, lib, params = relay.build(mod, target=build_target, params=params)

This runs into an error

TVMError: if is not supported.

Building this way does not support if statements in the graph

I am pretty sure this is caused by the if condition in

    if boxes.numel() == 0:
        return torch.empty((0,), dtype=torch.int64, device=boxes.device)
    # strategy: in order to perform NMS independently per class.
    # we add an offset to all the boxes. The offset is dependent
    # only on the class idx, and is large enough so that boxes
    # from different classes do not overlap
    else:
        max_coordinate = boxes.max()
        offsets = idxs.to(boxes) * (max_coordinate + torch.tensor(1).to(boxes))
        boxes_for_nms = boxes + offsets[:, None]
        keep = nms(boxes_for_nms, scores, iou_threshold)
        return keep

Do you have an idea how i can restructure the NMS avoiding this if statement?
All ideas are greatly appreciated!

[Tomakko]

Why does the torch tracing even record this if condiftion? As the docs say

Tracing only correctly records functions and modules which are not data dependent (e.g., do not have conditionals on data in tensors)

However when if print the mod object from

mod, params = relay.frontend.from_pytorch(scipt_module, shape_list)

I can see, that it includes an if condition, which i suspect to be the if condition from the NMS code snippet above.

 %500 = ndarray_size(%499, dtype="int32");
  %501 = equal(%500, 0);
  %538 = if (%501) {
    zeros(shape=[0], dtype="int64")
  } else {
    %502 = %494.1;
    %503 = cast(%502, dtype="int64");
    %504 = (%486, %497, %503);
    %505 = adv_index(%504);
    %506 = min(%505);
    %507 = subtract(%505, %506);
    %508 = add(%507, 1f);
    %509 = max(%499);
    %510 = cast(%503, dtype="float32");
    %511 = add(%509, 1f);
    %512 = multiply(%510, %511);
    %513 = expand_dims(%512, axis=1);
    %514 = expand_dims(%508, axis=-1);
    %515 = add(%499, %513);
    %516 = (%514, %515);
    %517 = concatenate(%516, axis=-1);
    %518 = shape_of(%508, dtype="int32");
    %519 = squeeze(%518);
    %520 = arange(0, %519, 1, start=meta[relay.Constant][24], stop=meta[relay.Call][0], step=meta[relay.Constant][25], dtype="int32");
    %521 = expand_dims(%517, axis=0);
    %522 = expand_dims(%520, axis=0);
    %523 = vision.non_max_suppression(%521, %518, %522, -1, 0.45f, meta[relay.attrs.NonMaximumSuppressionAttrs][0]);
    %524 = %523.0;
    %525 = cast_like(0, meta[relay.Constant][26]);
    %526 = squeeze(%524, axis=[0]);
    %527 = shape_of(%526, dtype="int32");
    %528 = cast_like(%527, meta[relay.Constant][26]);
    %529 = slice_like(%528, meta[relay.Constant][26], axes=None);
    %530 = less(meta[relay.Constant][26], %525);
    %531 = add(meta[relay.Constant][26], %529);
    %532 = where(%530, %531, meta[relay.Constant][26]);
    %533 = greater_equal(%532, %529);
    %534 = %523.1;
    %535 = where(%533, %529, %532);
    %536 = squeeze(%534, axis=[1]);
    %537 = dyn.strided_slice(%526, %535, %536, meta[relay.Constant][27], begin=None, end=None, strides=None, slice_mode="size", axes=None);
    cast(%537, dtype="int64")
  };

[zhiqwang]

Why does the torch tracing even record this if condition? As the docs say

Tracing only correctly records functions and modules which are not data dependent (e.g., do not have conditionals on data in tensors)

@Tomakko I have the same doubts as you.

The upstream torchvision team have refactored the batch_nms in the 0.10.0 release, seems that _batched_nms_vanilla has no dependence on if condition, I think that we could try this new version firstly.

from torchvision.ops import nms

@torch.jit._script_if_tracing
def batched_nms(
    boxes: Tensor,
    scores: Tensor,
    idxs: Tensor,
    iou_threshold: float,
) -> Tensor:
    # Based on Detectron2 implementation, just manually call nms() on each class independently
    keep_mask = torch.zeros_like(scores, dtype=torch.bool)
    for class_id in torch.unique(idxs):
        curr_indices = torch.where(idxs == class_id)[0]
        curr_keep_indices = nms(boxes[curr_indices], scores[curr_indices], iou_threshold)
        keep_mask[curr_indices[curr_keep_indices]] = True
    keep_indices = torch.where(keep_mask)[0]
    return keep_indices[scores[keep_indices].sort(descending=True)[1]]

Ref: https://github.com/pytorch/vision/blob/183a722169421c83638e68ee2d8fc5bd3415c4b4/torchvision/ops/boxes.py#L91-L105

I haven't dig deeper on tvm recently, let me know if this works for you.

[Tomakko]

Hi @zhiqwang ,

unfortunently it does not work. First, _batched_nms_vanilla is invoked after checking for numel with an if statement.

if boxes.numel() > 4_000 and not torchvision._is_tracing():
    return _batched_nms_vanilla(boxes, scores, idxs, iou_threshold)
else:
    return _batched_nms_coordinate_trick(boxes, scores, idxs, iou_threshold)

Second, multiple operations from _batched_nms_vanilla are not supported by TVM:

NotImplementedError: The following operators are not implemented: ['prim::unchecked_cast', 'prim::dtype', 'aten::unique_dim', 'aten::__isnot__']

[Tomakko]

Just FYI, i endet up tracing and exporting the model without the postprocessing which included the problematic nms operations. WIth this approach, i was able to compile the model with tvm.