make BoxSelector exportable to ONNX

Hi, I trained an object detecion model and I want to export it to ONNX with also the non max supression function.

Can you turn your monai.apps.detection.utils.box_selector.BoxSelector() class such as it works with torch.onnx.export()?

this is the error I get:

TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
elif self.dynamic or self.shape != shape:
[...]/python3.10/site-packages/monai/apps/detection/utils/box_selector.py:174: 
TracerWarning: Using len to get tensor shape might cause the trace to be incorrect. Recommended usage would be tensor.shape[0]. Passing a tensor of different shape might lead to errors or silently give incorrect results.
if len(boxes_list) != len(logits_list):
[...]/python3.10/site-packages/monai/apps/detection/utils/box_selector.py:187: TracerWarning: Iterating over a tensor might cause the trace to be incorrect. Passing a tensor of different shape won't change the number of iterations executed (and might lead to errors or silently give incorrect results).
for boxes_per_level, logits_per_level in zip(boxes_list, logits_list):
[...]/python3.10/site-packages/monai/apps/detection/utils/box_selector.py:137: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
num_topk = min(self.topk_candidates_per_level, flatten_topk_idxs.size(0))
[...]/python3.10/site-packages/monai/data/box_utils.py:409: TracerWarning: Converting a tensor to a Python integer might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
if int(boxes.shape[1] / 2) not in SUPPORTED_SPATIAL_DIMS:
[...]/python3.10/site-packages/monai/data/box_utils.py:413: TracerWarning: Converting a tensor to a Python integer might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
spatial_dims_set.add(int(boxes.shape[1] / 2))
[...]/python3.10/site-packages/monai/utils/type_conversion.py:141: TracerWarning: torch.as_tensor results are registered as constants in the trace. You can safely ignore this warning if you use this function to create tensors out of constant variables that would be the same every time you call this function. In any other case, this might cause the trace to be incorrect.
tensor = torch.as_tensor(tensor, **kwargs)
[...]/python3.10/site-packages/monai/utils/type_conversion.py:141: TracerWarning: Converting a tensor to a Python float might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
tensor = torch.as_tensor(tensor, **kwargs)
[...]/python3.10/site-packages/monai/data/box_utils.py:1171: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
if boxes.shape[0] == 0:
[...]/python3.10/site-packages/monai/data/box_utils.py:1101: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
if boxes.shape[0] == 0:
[...]/python3.10/site-packages/monai/data/box_utils.py:1104: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
if boxes.shape[0] != scores.shape[0]:
Traceback (most recent call last):
File "export.py", line 142, in <module>
torch.onnx.export(
File "[...]/python3.10/site-packages/torch/onnx/utils.py", line 516, in export
_export(
File "[...]/python3.10/site-packages/torch/onnx/utils.py", line 1613, in _export
graph, params_dict, torch_out = _model_to_graph(
File "[...]/python3.10/site-packages/torch/onnx/utils.py", line 1135, in _model_to_graph
graph, params, torch_out, module = _create_jit_graph(model, args)
File "[...]/python3.10/site-packages/torch/onnx/utils.py", line 1011, in _create_jit_graph
graph, torch_out = _trace_and_get_graph_from_model(model, args)
File "[...]/python3.10/site-packages/torch/onnx/utils.py", line 915, in _trace_and_get_graph_from_model
trace_graph, torch_out, inputs_states = torch.jit._get_trace_graph(
File "[...]/python3.10/site-packages/torch/jit/_trace.py", line 1296, in _get_trace_graph
outs = ONNXTracedModule(
File "[...]/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "[...]/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
File "[...]/python3.10/site-packages/torch/jit/_trace.py", line 138, in forward
graph, out = torch._C._create_graph_by_tracing(
File "[...]/python3.10/site-packages/torch/jit/_trace.py", line 129, in wrapper
outs.append(self.inner(*trace_inputs))
File "[...]/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "[...]/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
File "[...]/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _slow_forward
result = self.forward(*input, **kwargs)
File "/export.py", line 58, in forward
selected_boxes, selected_scores, selected_labels = self.nms.select_boxes_per_image(boxes_list, logits_list, 
(im.size(2), im.size(3)))
File "[...]/python3.10/site-packages/monai/apps/detection/utils/box_selector.py", line 206, in select_boxes_per_image
keep_t: Tensor = batched_nms(  # type: ignore
File "[...]/python3.10/site-packages/monai/data/box_utils.py", line 1186, in batched_nms
keep = non_max_suppression(boxes_for_nms, scores_t, nms_thresh, max_proposals, box_overlap_metric)
File "[...]/python3.10/site-packages/monai/data/box_utils.py", line 1115, in non_max_suppression
boxes_sort = deepcopy(boxes_t)[sort_idxs, :]
File "[...]/python3.10/copy.py", line 153, in deepcopy
y = copier(memo)
File "[...]/python3.10/site-packages/torch/_tensor.py", line 122, in __deepcopy__
new_storage = self._typed_storage()._deepcopy(memo)
File "[...]/python3.10/site-packages/torch/storage.py", line 839, in _deepcopy
return self._new_wrapped_storage(copy.deepcopy(self._untyped_storage, memo))
File "[...]/python3.10/copy.py", line 153, in deepcopy
y = copier(memo)
File "[...]/python3.10/site-packages/torch/storage.py", line 112, in __deepcopy__
new_storage = self.clone()
File "[...]/python3.10/site-packages/torch/storage.py", line 126, in clone
return type(self)(self.nbytes(), device=self.device).copy_(self)
RuntimeError: NYI: Named tensors are not supported with the tracer

Thanks in advance.

Start with monai/apps/detection/utils/box_selector.py, especially BoxSelector.select_boxes_per_image, and trace its calls into monai/data/box_utils.py. Reproduce the failure through torch.onnx.export using the reported model path. Done means BoxSelector, including non-maximum suppression, exports without the reported tracer or named-tensor errors and remains usable for the intended inputs.