Cannot export detection model to onnx with dynamo option

[Rayndell]

🐛 Describe the bug

I am trying to export a PyTorch (faster rcnn with resnet 101 backbone) to onnx. The static export (with dynamo=False) works well to a certain extent (some transforms v2 ops are not fully exportable for instance), but I'm struggling to do the same with dynamo=True, which is as I understood the proper way to do it. Here is my sample code:

import torch
from torchvision.tv_tensors import BoundingBoxes
import onnx
import torch.nn as nn
from torch.export.dynamic_shapes import Dim
import configparser
from model import create_model
import argparse
from torchvision.transforms import v2
import torchvision.ops as ops
from torchvision.models.detection.rpn import AnchorGenerator
from torchvision.models.detection.backbone_utils import resnet_fpn_backbone
rom torchvision.models.detection import FasterRCNN
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor

def create_model(num_classes, load_weights=True):
    detections_per_img = 1000
    score_thresh = 0.6
    nms_thresh = 0.6
    resize_height = 1216
    resize_width = 1936

    # Use a custom backbone
    backbone_name = "resnet101"
    backbone = resnet_fpn_backbone(backbone_name=backbone_name, weights='DEFAULT' if load_weights else None)
        
    # Modifying bounding boxes candidates generation
    anchor_sizes = ((32,), (64,), (128,), (256,), (512,),)
    aspect_ratios=((0.05, 0.1, 0.2, 0.33, 0.5, 1.0, 2.0, 3.0),) * len(anchor_sizes)
    anchor_generator = AnchorGenerator(sizes=anchor_sizes, aspect_ratios=aspect_ratios)
    roi_pooler = torchvision.ops.MultiScaleRoIAlign(
                featmap_names=['0', '1', '2', '3', 'pool'],
                output_size=7,
                sampling_ratio=2
            )
        
    model = FasterRCNN(backbone=backbone, num_classes=num_classes,
                rpn_anchor_generator=anchor_generator, box_roi_pool=roi_pooler)

    model.roi_heads.detections_per_img = detections_per_img
    model.roi_heads.score_thresh = score_thresh

    # Modifying the last layers to get the right number of classes
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    return model

class DetectionModel(nn.Module):
    def __init__(self, base_model, resize_height, resize_width, nms_threshold=0.6):
        super(DetectionModel, self).__init__()
        self.base_model = base_model
        self.nms_threshold = nms_threshold
        self.resize_height = resize_height
        self.resize_width = resize_width
        self.transform = v2.Compose([
            v2.ToImage(),
            v2.Resize(size=(resize_height, resize_width), antialias=False),
            v2.ToDtype(torch.float32, scale=True)
        ])
        
    def forward(self, input):
        img_height, img_width = input.shape[-2:]
        resize = v2.Resize(size=(img_height, img_width), antialias=False)
        
        # Apply the preprocessing steps
        input = self.transform(input)
        
        # Apply the base model to the preprocessed input
        output = self.base_model(input)
        
        all_boxes = []
        all_scores = []
        all_labels = []
        
        for i in range(input.shape[0]):
            pred = output[i]
            boxes = pred['boxes']
            scores = pred['scores']
            labels = pred['labels']

            # Non max suppression
            keep = ops.nms(boxes, scores, iou_threshold=self.nms_threshold)
            boxes = boxes[keep]
            scores = scores[keep]
            labels = labels[keep]
            
            # Replacing the boxes coordinates in the original image before resizing
            boxes = BoundingBoxes(boxes, format="XYXY", canvas_size=(self.resize_height, self.resize_width))
            boxes = resize(boxes)
            
            all_boxes.append(boxes)
            all_scores.append(scores)
            all_labels.append(labels)
        
        return torch.stack(all_boxes), torch.stack(all_scores), torch.stack(all_labels)
    
    def export_to_onnx(self, path, device):
        print("Exporting model to ONNX...")
        self.eval()
        input_names = ["input"]
        output_names = ["boxes", "scores", "labels"]
        
        input_example = torch.randint(0, 255, (1, 3, self.resize_height, self.resize_width), dtype=torch.uint8).to(device)

        dynamic_shapes={
            "input": {0: "batch_size", 2: "height", 3: "width"},
        }
        onnx_program = torch.onnx.export(
            self, input_example,
            dynamic_shapes=dynamic_shapes,
            input_names=input_names, output_names=output_names, dynamo=True)
        
        onnx_program.save(path)
        
        print(f"ONNX model exported to {path}")

def main():
    resize_height = 1216
    resize_width = 1936
    nms_threshold = 0.6
    
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
    # we load the given checkpoint
    checkpoint = torch.load("checkpoint_epoch_50.pth", map_location=device, weights_only=False)
    
    # Reconstruct the model
    base_model = create_model(num_classes=70, load_weights=False)
    base_model.load_state_dict(checkpoint["model_state_dict"], strict=False)
    base_model = base_model.to(device)
    model = DetectionModel(base_model, resize_height, resize_width, nms_threshold)
    
    model.export_to_onnx("model.onnx", device)

if __name__ == "__main__":
    main()

To run this you will need the checkpoint, it is downloadable here: https://evolucare-my.sharepoint.com/:u:/p/a_ducournau/IQB0b8lRnuCET52fCepH4ANlAU_LbohDImHcvraRYSZiaCw?e=dHjmtV

Here is the detail of the error messages I received when execution torch.onn.export:

Exporting model to ONNX...
[torch.onnx] Obtain model graph for DetectionModel([...] with torch.export.export(..., strict=False)...
[torch.onnx] Obtain model graph for DetectionModel([...] with torch.export.export(..., strict=False)... ❌
[torch.onnx] Obtain model graph for DetectionModel([...] with torch.export.export(..., strict=True)...
[torch.onnx] Obtain model graph for DetectionModel([...] with torch.export.export(..., strict=True)... ❌
[torch.onnx] Obtain model graph for DetectionModel([...] with torch.export draft_export...
[torch.onnx] Obtain model graph for DetectionModel([...] with torch.export draft_export... ❌
Traceback (most recent call last):
File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 1809, in path_of_module
return Tracer.path_of_module(self, mod)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 482, in path_of_module
raise NameError("module is not installed as a submodule")
NameError: module is not installed as a submodule

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 1905, in call_module
return Tracer.call_module(self, m, forward, args, kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 527, in call_module
module_qualified_name = self.path_of_module(m)
^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 1811, in path_of_module
raise _ModuleNotInstalledAsSubmoduleError from e
torch.fx.experimental.proxy_tensor._ModuleNotInstalledAsSubmoduleError

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "/usr/local/lib/python3.12/dist-packages/torch/onnx/_internal/exporter/_capture_strategies.py", line 118, in call
exported_program = self._capture(model, args, kwargs, dynamic_shapes)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/onnx/_internal/exporter/_capture_strategies.py", line 210, in _capture
return torch.export.export(
^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/export/init.py", line 319, in export
raise e
File "/usr/local/lib/python3.12/dist-packages/torch/export/init.py", line 286, in export
return _export(
^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/export/_trace.py", line 1164, in wrapper
raise e
File "/usr/local/lib/python3.12/dist-packages/torch/export/_trace.py", line 1130, in wrapper
ep = fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/export/exported_program.py", line 123, in wrapper
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/export/_trace.py", line 2176, in _export
ep = _export_for_training(
^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/export/_trace.py", line 1164, in wrapper
raise e
File "/usr/local/lib/python3.12/dist-packages/torch/export/_trace.py", line 1130, in wrapper
ep = fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/export/exported_program.py", line 123, in wrapper
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/export/_trace.py", line 2037, in _export_for_training
export_artifact = export_func(
^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/export/_trace.py", line 1979, in _non_strict_export
aten_export_artifact = _to_aten_func( # type: ignore[operator]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/export/_trace.py", line 1770, in _export_to_aten_ir_make_fx
gm, graph_signature = transform(_make_fx_helper)(
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/export/_trace.py", line 1900, in _aot_export_non_strict
gm, sig = aot_export(wrapped_mod, args, kwargs=kwargs, **flags)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/export/_trace.py", line 1685, in _make_fx_helper
gm = make_fx(
^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 2318, in wrapped
return make_fx_tracer.trace(f, *args)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 2250, in trace
return self._trace_inner(f, *args)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 2221, in _trace_inner
t = dispatch_trace(
^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/_compile.py", line 53, in inner
return disable_fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py", line 929, in _fn
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 1254, in dispatch_trace
graph = tracer.trace(root, concrete_args) # type: ignore[arg-type]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 1835, in trace
res = super().trace(root, concrete_args)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 850, in trace
(self.create_arg(fn(*args)),),
^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 1312, in wrapped
out = f(*tensors) # type:ignore[call-arg]
^^^^^^^^^^^
File "", line 1, in
File "/usr/local/lib/python3.12/dist-packages/torch/export/_trace.py", line 1589, in wrapped_fn
return tuple(flat_fn(*args))
^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/utils.py", line 184, in flat_fn
tree_out = fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 906, in functional_call
out = mod(*args[params_len:], **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 825, in module_call_wrapper
return self.call_module(mod, forward, args, kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 1905, in call_module
return Tracer.call_module(self, m, forward, args, kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 542, in call_module
ret_val = forward(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 818, in forward
return _orig_module_call(mod, *args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1784, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/export/_trace.py", line 1884, in forward
tree_out = mod(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 825, in module_call_wrapper
return self.call_module(mod, forward, args, kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 1905, in call_module
return Tracer.call_module(self, m, forward, args, kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 542, in call_module
ret_val = forward(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 818, in forward
return _orig_module_call(mod, *args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1784, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/workspace/detection/export_onnx.py", line 49, in forward
input = self.transform(input)
^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 825, in module_call_wrapper
return self.call_module(mod, forward, args, kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 1905, in call_module
return Tracer.call_module(self, m, forward, args, kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 542, in call_module
ret_val = forward(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 818, in forward
return _orig_module_call(mod, *args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1784, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torchvision/transforms/v2/_container.py", line 52, in forward
outputs = transform(*inputs)
^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 825, in module_call_wrapper
return self.call_module(mod, forward, args, kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 1914, in call_module
return forward(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/_symbolic_trace.py", line 818, in forward
return _orig_module_call(mod, *args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1784, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torchvision/transforms/v2/_transform.py", line 69, in forward
self.transform(inpt, params) if needs_transform else inpt
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torchvision/transforms/v2/_type_conversion.py", line 39, in transform
return F.to_image(inpt)
^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torchvision/transforms/v2/functional/_type_conversion.py", line 23, in to_image
return tv_tensors.Image(output)
^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torchvision/tv_tensors/_image.py", line 50, in new
return tensor.as_subclass(cls)
^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/utils/_stats.py", line 28, in wrapper
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 872, in torch_dispatch
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 829, in call
return self._op(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 1360, in torch_function
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 829, in call
return self._op(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 1407, in torch_function
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 829, in call
return self._op(*args, *kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/_export/non_strict_utils.py", line 1051, in torch_function
return func(args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 829, in call
return self._op(args, *kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/utils/_stats.py", line 28, in wrapper
return fn(args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 1352, in torch_dispatch
return self.dispatch(func, types, args, kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 2058, in dispatch
return self._cached_dispatch_impl(func, types, args, kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 1509, in _cached_dispatch_impl
entry = self._make_cache_entry(state, key, func, args, kwargs, output)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 1878, in _make_cache_entry
output_info = self._get_output_info_for_cache_entry(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 1794, in _get_output_info_for_cache_entry
synth_output = self._output_from_cache_entry(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 1979, in _output_from_cache_entry
return self._get_output_tensor_from_cache_entry(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 1954, in get_output_tensor_from_cache_entry
empty.set(storage, storage_offset, shape, stride)
File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/sym_node.py", line 381, in gt
return self._gt(other) # type: ignore[attr-defined]
^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/sym_node.py", line 1295, in wrapper
result = fn(self, other)
^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/sym_node.py", line 1350, in binary_magic_impl
self, handle_sym_dispatch(op, (wrap_node(self), wrap_node(other)), {})
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 2366, in handle_sym_dispatch
return mode.sym_dispatch(func, types, args, kwargs) # type: ignore[arg-type, return-value]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 1548, in sym_dispatch
p_out_thunk = thunkify(
^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 508, in thunkify
r = f(args, **kwargs)
^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 1506, in _compute_proxy
n_args = tuple(
^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 1508, in
get_proxy_slot(a, self.tracer).force().node
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 348, in get_proxy_slot
raise RuntimeError(
RuntimeError: 3s16s82(s31 - 1) + 2s16s82 + s82(s16 - 1) + s82 (140038440703120)is not tracked with proxy for <torch.fx.experimental.proxy_tensor._ModuleStackTracer object at 0x7f5d3dc808c0>

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/workspace/detection/export_onnx.py", line 179, in
main(args)
File "/workspace/detection/export_onnx.py", line 170, in main
model.export_to_onnx(args.output, device)
File "/workspace/detection/export_onnx.py", line 109, in export_to_onnx
onnx_program = torch.onnx.export(
^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/onnx/init.py", line 377, in export
return _compat.export_compat(
^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/onnx/_internal/exporter/_compat.py", line 121, in export_compat
onnx_program = _core.export(
^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/onnx/_internal/exporter/_flags.py", line 20, in wrapper
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/onnx/_internal/exporter/_core.py", line 1379, in export
raise _errors.TorchExportError(
torch.onnx._internal.exporter._errors.TorchExportError: Failed to export the model with torch.export. This is step 1/3 of exporting the model to ONNX. Next steps:

• Modify the model code for torch.export.export to succeed. Refer to https://pytorch.org/docs/stable/generated/exportdb/index.html for more information.
• Debug torch.export.export and summit a PR to PyTorch.
• Create an issue in the PyTorch GitHub repository against the torch.export component and attach the full error stack as well as reproduction scripts.

Exception summary

<class 'RuntimeError'>: 3s16s82*(s31 - 1) + 2s16s82 + s82*(s16 - 1) + s82 (140038440703120)is not tracked with proxy for <torch.fx.experimental.proxy_tensor._ModuleStackTracer object at 0x7f5d3dc808c0>

(Refer to the full stack trace above for more information.)

Versions

Collecting environment information...
PyTorch version: 2.8.0+cu126
Is debug build: False
CUDA used to build PyTorch: 12.6
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.5 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.35

Python version: 3.12.13 (main, Mar 4 2026, 09:23:07) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-6.1.0-40-amd64-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.6.85
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA RTX 5000 Ada Generation
Nvidia driver version: 575.57.08
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.9.5.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv.so.9.5.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn.so.9.5.1
/usr/lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.5.1
/usr/lib/x86_64-linux-gnu/libcudnn_engines_runtime_compiled.so.9.5.1
/usr/lib/x86_64-linux-gnu/libcudnn_graph.so.9.5.1
/usr/lib/x86_64-linux-gnu/libcudnn_heuristic.so.9.5.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops.so.9.5.1
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True
Caching allocator config: N/A

CPU:
Architecture: x86_64
CPU op-mode(s): 32-bit, 64-bit
Address sizes: 52 bits physical, 57 bits virtual
Byte Order: Little Endian
CPU(s): 24
On-line CPU(s) list: 0-23
Vendor ID: GenuineIntel
Model name: Intel(R) Xeon(R) w5-3423
CPU family: 6
Model: 143
Thread(s) per core: 2
Core(s) per socket: 12
Socket(s): 1
Stepping: 8
CPU max MHz: 4200.0000
CPU min MHz: 800.0000
BogoMIPS: 4224.00
Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l3 cat_l2 cdp_l3 invpcid_single intel_ppin cdp_l2 ssbd mba ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb intel_pt avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local split_lock_detect avx_vnni avx512_bf16 wbnoinvd dtherm ida arat pln pts hwp hwp_act_window hwp_epp hwp_pkg_req avx512vbmi umip pku ospke waitpkg avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg tme avx512_vpopcntdq la57 rdpid bus_lock_detect cldemote movdiri movdir64b enqcmd fsrm md_clear serialize tsxldtrk pconfig arch_lbr ibt amx_bf16 avx512_fp16 amx_tile amx_int8 flush_l1d arch_capabilities ibpb_exit_to_user
Virtualization: VT-x
L1d cache: 576 KiB (12 instances)
L1i cache: 384 KiB (12 instances)
L2 cache: 24 MiB (12 instances)
L3 cache: 30 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-23
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2: Mitigation; Enhanced / Automatic IBRS; IBPB conditional; PBRSB-eIBRS SW sequence; BHI BHI_DIS_S
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Mitigation; IBPB before exit to userspace

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] nvidia-cublas-cu12==12.6.4.1
[pip3] nvidia-cuda-cupti-cu12==12.6.80
[pip3] nvidia-cuda-nvrtc-cu12==12.6.77
[pip3] nvidia-cuda-runtime-cu12==12.6.77
[pip3] nvidia-cudnn-cu12==9.10.2.21
[pip3] nvidia-cufft-cu12==11.3.0.4
[pip3] nvidia-curand-cu12==10.3.7.77
[pip3] nvidia-cusolver-cu12==11.7.1.2
[pip3] nvidia-cusparse-cu12==12.5.4.2
[pip3] nvidia-cusparselt-cu12==0.7.1
[pip3] nvidia-nccl-cu12==2.27.3
[pip3] nvidia-nvjitlink-cu12==12.6.85
[pip3] nvidia-nvtx-cu12==12.6.77
[pip3] onnx==1.20.1
[pip3] onnx-ir==0.1.15
[pip3] onnxruntime==1.23.2
[pip3] onnxscript==0.5.7
[pip3] torch==2.8.0+cu126
[pip3] torchaudio==2.8.0+cu126
[pip3] torchmetrics==1.8.2
[pip3] torchsummary==1.5.1
[pip3] torchvision==0.23.0+cu126
[pip3] triton==3.4.0
[conda] Could not collect

[Rayndell]

Hi, I was wondering it there was something to do with the above issue. I am trying to do the same with a more standard classification model (resnet101) and the error is the same:

NameError: module is not installed as a submodule

Does it come from the preprocessing functions I want to include before the call to the model?

[zy1git]

@Rayndell Thanks for raising this issue!

Including v2 transforms inside the model's forward() is currently not compatible with dynamo-based ONNX export. As a workaround, you can move the transforms outside the model as a separate preprocessing step and export only the base model.

Does it come from the preprocessing functions I want to include before the call to the model?

Yes.