Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Enhance Error Handling for TENSOR_META in Output Shape Calculation for YOLOv8 Quantization #1186

Merged
merged 1 commit into from
Sep 1, 2024

Conversation

ambitious-octopus
Copy link
Contributor

@ambitious-octopus ambitious-octopus commented Aug 22, 2024

Pull Request Description:

The PR addresses the case where a PyTorch graph contains additional input nodes, like args and kwargs, that have no outputs. Instead of causing an error, these nodes will be removed in the subsequent method: remove_broken_nodes_from_graph.

This PR includes a small update to the code that handles the extraction of output_shape based on the type of the _node.meta[TYPE].

The old code assumes that _node.meta[TENSOR_META] is always present when _node.meta[TYPE] is a list or tuple. This assumption could lead to a KeyError if TENSOR_META is missing from _node.meta.

The new code uses _node.meta.get(TENSOR_META, []) instead, which returns an empty list if TENSOR_META is not found, thereby preventing potential errors and ensuring the code handles such cases more gracefully.

This change is necessary to handle scenarios specific to YOLOv8 quantization, where TENSOR_META may not always be present in _node.meta.

Checklist before requesting a review:

  • I set the appropriate labels on the pull request.
  • I have added/updated the release note draft (if necessary).
  • I have updated the documentation to reflect my changes (if necessary).
  • All function and files are well documented.
  • All function and classes have type hints.
  • There is a licenses in all file.
  • The function and variable names are informative.
  • I have checked for code duplications.
  • I have added new unittest (if necessary).

Co-authored-by: Laughing-q <1185102784@qq.com>
@Idan-BenAmi
Copy link
Collaborator

Hi @ambitious-octopus , thanks for your PR, can you explain the specific case that reproduce this issue?

@ambitious-octopus
Copy link
Contributor Author

Hey @Idan-BenAmi, you can reproduce the error with this script. This script closely resembles the notebook tutorial located at tutorials/notebooks/imx500_notebooks/pytorch/pytorch_yolov8n_for_imx500.ipynb. However, instead of using the custom-modified YOLOv8n, it utilizes the original YOLOv8n model from the Ultralytics package with minor adjustments for compatibility with torch.fx. To install this specific version, you can refer to this pull request or simply install it directly via pip:

pip install git+https://github.com/ultralytics/ultralytics.git@quan

Running the script without this fix results in an error when invoking the mct.core.pytorch_resource_utilization_data function.

Anonymized output
Traceback (most recent call last):
  File "/home/me/repos/ultralytics/test.py", line 3, in <module>
    model.export(format="mct", batch=2)
  File "/home/me/repos/ultralytics/ultralytics/engine/model.py", line 744, in export
    return Exporter(overrides=args, _callbacks=self.callbacks)(model=self.model)
  File "/home/me/miniconda3/envs/ultralytics/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/me/repos/ultralytics/ultralytics/engine/exporter.py", line 331, in __call__
    f[12], _ = self.export_mct()
  File "/home/me/repos/ultralytics/ultralytics/engine/exporter.py", line 143, in outer_func
    raise e
  File "/home/me/repos/ultralytics/ultralytics/engine/exporter.py", line 138, in outer_func
    f, model = inner_func(*args, **kwargs)
  File "/home/me/repos/ultralytics/ultralytics/engine/exporter.py", line 1077, in export_mct
    resource_utilization_data = mct.core.pytorch_resource_utilization_data(in_model=self.model,
  File "/home/me/repos/model_optimization/model_compression_toolkit/core/pytorch/resource_utilization_data_facade.py", line 83, in pytorch_resource_utilization_data
    return compute_resource_utilization_data(in_model,
  File "/home/me/repos/model_optimization/model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization_data.py", line 64, in compute_resource_utilization_data
    transformed_graph = graph_preparation_runner(in_model,
  File "/home/me/repos/model_optimization/model_compression_toolkit/core/graph_prep_runner.py", line 72, in graph_preparation_runner
    graph = read_model_to_graph(in_model,
  File "/home/me/repos/model_optimization/model_compression_toolkit/core/graph_prep_runner.py", line 207, in read_model_to_graph
    graph = fw_impl.model_reader(in_model,
  File "/home/me/repos/model_optimization/model_compression_toolkit/core/pytorch/pytorch_implementation.py", line 149, in model_reader
    return model_reader(_module, representative_data_gen, self.to_numpy, self.to_tensor)
  File "/home/me/repos/model_optimization/model_compression_toolkit/core/pytorch/reader/reader.py", line 154, in model_reader
    graph = build_graph(fx_model, to_numpy)
  File "/home/me/repos/model_optimization/model_compression_toolkit/core/pytorch/reader/reader.py", line 65, in build_graph
    nodes, inputs, outputs, fx_node_2_graph_node = nodes_builder(model, module_dict, to_numpy)
  File "/home/me/repos/model_optimization/model_compression_toolkit/core/pytorch/reader/graph_builders.py", line 238, in nodes_builder
    input_shape, output_shape = _extract_input_and_output_shapes(node)
  File "/home/me/repos/model_optimization/model_compression_toolkit/core/pytorch/reader/graph_builders.py", line 151, in _extract_input_and_output_shapes
    output_shape = [list(m.shape) for m in _node.meta[TENSOR_META]]
KeyError: 'tensor_meta'

@Idan-BenAmi
Copy link
Collaborator

Idan-BenAmi commented Aug 26, 2024

Hi @ambitious-octopus, I'm trying to reproduce the issue using the branch you sent to me (quan), both via pip (as you mentioned) and also by cloning the ultralytics repository and switching quan branch.
I got a torch FX error (see below), which it seems you resolved already.

Can you advise? am I using the correct branch/version?
Thanks

2024-08-26 13:25:24 idanb-ub20 Model Compression Toolkit[2611824] CRITICAL Error parsing model with torch.fx
fx error: symbolically traced variables cannot be used as inputs to control flow
Traceback (most recent call last):
  File "/data/projects/swat/users/idanb/repository/git4/model_optimization/model_compression_toolkit/core/pytorch/reader/reader.py", line 90, in fx_graph_module_generation
    symbolic_traced = symbolic_trace(pytorch_model)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/projects/swat/envs/idanb/venv_np_1_26/lib/python3.11/site-packages/torch/fx/_symbolic_trace.py", line 1193, in symbolic_trace
    graph = tracer.trace(root, concrete_args)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/projects/swat/envs/idanb/venv_np_1_26/lib/python3.11/site-packages/torch/fx/_symbolic_trace.py", line 793, in trace
    (self.create_arg(fn(*args)),),
                     ^^^^^^^^^
  File "/data/projects/swat/users/idanb/repository/git4/ultralytics/ultralytics/nn/tasks.py", line 0, in wrapper
  File "/data/projects/swat/users/idanb/repository/git4/ultralytics/ultralytics/nn/tasks.py", line 129, in forward
    return self.predict(x, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/projects/swat/users/idanb/repository/git4/ultralytics/ultralytics/nn/tasks.py", line 147, in predict
    return self._predict_once(x, profile, visualize, embed)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/projects/swat/users/idanb/repository/git4/ultralytics/ultralytics/nn/tasks.py", line 168, in _predict_once
    x = m(x)  # run
        ^^^^
  File "/data/projects/swat/envs/idanb/venv_np_1_26/lib/python3.11/site-packages/torch/fx/_symbolic_trace.py", line 771, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/projects/swat/envs/idanb/venv_np_1_26/lib/python3.11/site-packages/torch/fx/_symbolic_trace.py", line 495, in call_module
    ret_val = forward(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/projects/swat/envs/idanb/venv_np_1_26/lib/python3.11/site-packages/torch/fx/_symbolic_trace.py", line 764, in forward
    return _orig_module_call(mod, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/projects/swat/envs/idanb/venv_np_1_26/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/projects/swat/envs/idanb/venv_np_1_26/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/projects/swat/users/idanb/repository/git4/ultralytics/ultralytics/nn/modules/head.py", line 60, in forward
    y = self._inference(x)
        ^^^^^^^^^^^^^^^^^^
  File "/data/projects/swat/users/idanb/repository/git4/ultralytics/ultralytics/nn/modules/head.py", line 98, in _inference
    elif self.dynamic or self.shape != shape:
                         ^^^^^^^^^^^^^^^^^^^
  File "/data/projects/swat/envs/idanb/venv_np_1_26/lib/python3.11/site-packages/torch/fx/proxy.py", line 443, in __bool__
    return self.tracer.to_bool(self)
           ^^^^^^^^^^^^^^^^^^^^^^^^^

@Idan-BenAmi Idan-BenAmi merged commit 11521f5 into sony:main Sep 1, 2024
29 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants