No module named 'onnx.optimizer'

[giaxxi]

I'm using this script to convert a model: https://github.com/facebookresearch/detectron2/blob/f5e4c452bba854b8ed14a7240f51720ca7857e91/tools/deploy/export_model.py#L220

For example using:

export_method = "caffe2_tracing" # "caffe2_tracing", "scripting" or "tracing"
export_format = "caffe2" # "caffe2", "onnx" or "torchscript"

I get the following error:

~/.local/lib/python3.8/site-packages/caffe2/python/onnx/backend.py in <module>
     26 import onnx.numpy_helper
     27 import onnx.defs
---> 28 import onnx.optimizer
     29 import onnx.shape_inference
     30 import onnx.utils

ModuleNotFoundError: No module named 'onnx.optimizer'

The installed ONNX version is the latest:
(Successfully installed onnx-1.10.1)

onnx.__version__
#=> '1.10.1'

Even if I import ONNX in a fresh Colab and try to import onnx.optimizer I get the same error.

Any way to fix it?
Thank you.

[jcwchen]

Hi @giaxxi,
ONNX has moved optimizer to another repo since 1.9.0: https://github.com/onnx/onnx/releases/tag/v1.9.0

Which PyTorch are you using? You should not bump into this error with the latest PyTorch because it seems they has already handled it: pytorch/pytorch#55889. If you are using an old version of PyTorch, using ONNX <= 1.8.1 would be a quick workaround. But, I would suggest you use the latest PyTorch with latest optimizer from onnx/optimizer repo because the new optimizer has been much improved from the old one. Thanks.

[jcwchen]

It should be fixed by pytorch/pytorch#75718. Next PyTorch release you should be able to use onnxoptimizer in caffe2. Thanks for catching this issue.

[bouachalazhar]

It should be fixed by pytorch/pytorch#75718. Next PyTorch release you should be able to use onnxoptimizer in caffe2. Thanks for catching this issue.

Unfornately, I have the same problem with PyTorch : 1.13.1+cpu, ONNX : 1.13.1, detectron2 : 0.6

[jcwchen]

Thanks for reporting. Could you please provide the reproduction steps and the error message you saw?

[bouachalazhar]

Thanks for reporting. Could you please provide the reproduction steps and the error message you saw?

Of course,

import warnings
from pathlib import Path, PurePosixPath

from IPython.display import Markdown, display
from openvino.runtime import Core

from detectron2.export import dump_torchscript_IR, torchscript
from detectron2.modeling import build_model
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.export.caffe2_export import export_onnx_model
import onnx
import onnxruntime as rt
import detectron2.data.transforms as T

def mask_image(im, outputs):
    v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    return out.get_image()[:, :, ::-1]

image_filename = "data/test.jpg"
image = cv2.imread(image_filename)

IMAGE_WIDTH, IMAGE_HEIGHT = image.shape[1::-1]
DIRECTORY_NAME = "models"

cfg = get_cfg()
if torch.cuda.is_available():
    cfg.MODEL.DEVICE = 'cuda'
else:
    cfg.MODEL.DEVICE = 'cpu'
# add project-specific config (e.g., TensorMask) here if you're not running a model in detectron2's core library
path = PurePosixPath("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
cfg.merge_from_file(model_zoo.get_config_file(path))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # set threshold for this model
# Find a model from detectron2's model zoo. You can use the https://dl.fbaipublicfiles... url as well
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(str(path))
# predictor = DefaultPredictor(cfg)
# Build model and prepare input
model = build_model(cfg)
model.eval()
checkpointer = DetectionCheckpointer(model)
checkpointer.load(cfg.MODEL.WEIGHTS)
aug = T.ResizeShortestEdge([cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST], cfg.INPUT.MAX_SIZE_TEST)
height, width = image.shape[:2]
image = aug.get_transform(image).apply_image(image)
image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))
inputs = {"image": image, "height": height, "width": width}

BASE_MODEL_NAME = DIRECTORY_NAME + "/" + path.stem
weights_path = Path(BASE_MODEL_NAME + ".pt")

# Paths where ONNX and OpenVINO IR models will be stored.
onnx_path = weights_path.with_suffix('.onnx')
if not onnx_path.parent.exists():
    onnx_path.parent.mkdir()
ir_path = onnx_path.with_suffix(".xml")

# print("Downloading the {} model (if it has not been downloaded already)...".format(BASE_MODEL_NAME.split("/")[1]))
# weights = getattr(torchvision.models.segmentation, weights)
# download_file(weights.COCO_WITH_VOC_LABELS_V1.url, filename=weights_path.name, directory=weights_path.parent)
# # create model object
# model = getattr(torchvision.models.segmentation, model_name)()
# # read state dict, use map_location argument to avoid a situation where weights are saved in cuda (which may not be unavailable on the system)
# state_dict = torch.load(weights_path, map_location='cpu')
# if BASE_MODEL_NAME.split("/")[1] == "deeplabv3_resnet50":
#     for key in list(state_dict.keys())[-8:]:
#         del state_dict[key]
# # load state dict to model
# model.load_state_dict(state_dict)
# # switch model from training to inference mode
# model.eval()
# print("Loaded PyTorch {} model".format(BASE_MODEL_NAME.split("/")[1]))


# Export to Onnx model
onnxModel = export_onnx_model(cfg, model, [inputs])
onnx.save(onnxModel, onnx_path)

# with warnings.catch_warnings():
#     warnings.filterwarnings("ignore")
#     if not onnx_path.exists():
#         dummy_input = torch.randn(1, 3, IMAGE_HEIGHT, IMAGE_WIDTH)
#         torch.onnx.export(
#             model,
#             dummy_input,
#             onnx_path,
#         )
#         print(f"ONNX model exported to {onnx_path}.")
#     else:
#         print(f"ONNX model {onnx_path} already exists.")

# Construct the command for Model Optimizer.
mo_command = f"""mo
                --input_model "{onnx_path}"
                --compress_to_fp16
                --output_dir "{ir_path.parent}"
                """
mo_command = " ".join(mo_command.split())
print("Model Optimizer command to convert the ONNX model to OpenVINO:")
display(Markdown(f"`{mo_command}`"))

if not ir_path.exists():
    print("Exporting ONNX model to IR... This may take a few minutes.")
    mo_result = %sx $mo_command
    print("\n".join(mo_result))
else:
    print(f"IR model {ir_path} already exists.")


# Load the network to OpenVINO Runtime.
ie = Core()
model_onnx = ie.read_model(model=onnx_path)
compiled_model_onnx = ie.compile_model(model=model_onnx, device_name="CPU")

output_layer_onnx = compiled_model_onnx.output(0)

# Run inference on the input image.
res_onnx = compiled_model_onnx([image])[output_layer_onnx]

# Convert the network result to a segmentation map and display the result.
result_mask_onnx = predictor(image)
cv2.imshow(mask_image(image, result_mask_onnx))

# Load the network in OpenVINO Runtime.
ie = Core()
model_ir = ie.read_model(model=ir_path)
compiled_model_ir = ie.compile_model(model=model_ir, device_name="CPU")

# Get input and output layers.
output_layer_ir = compiled_model_ir.output(0)

# Run inference on the input image.
res_ir = compiled_model_ir([image])[output_layer_ir]

result_mask_ir = predictor(image)
cv2.imshow(mask_image(image, result_mask_ir))

model.eval()
with torch.no_grad():
    result_torch = model(image)

result_mask_torch = predictor(image)
cv2.imshow(mask_image(image, result_mask_torch))

num_images = 100

with torch.no_grad():
    start = time.perf_counter()
    for _ in range(num_images):
        model(image)
    end = time.perf_counter()
    time_torch = end - start
print(
    f"PyTorch model on CPU: {time_torch/num_images:.3f} seconds per image, "
    f"FPS: {num_images/time_torch:.2f}"
)

start = time.perf_counter()
for _ in range(num_images):
    compiled_model_onnx([image])
end = time.perf_counter()
time_onnx = end - start
print(
    f"ONNX model in OpenVINO Runtime/CPU: {time_onnx/num_images:.3f} "
    f"seconds per image, FPS: {num_images/time_onnx:.2f}"
)

start = time.perf_counter()
for _ in range(num_images):
    compiled_model_ir([image])
end = time.perf_counter()
time_ir = end - start
print(
    f"OpenVINO IR model in OpenVINO Runtime/CPU: {time_ir/num_images:.3f} "
    f"seconds per image, FPS: {num_images/time_ir:.2f}"
)

if "GPU" in ie.available_devices:
    compiled_model_onnx_gpu = ie.compile_model(model=model_onnx, device_name="GPU")
    start = time.perf_counter()
    for _ in range(num_images):
        compiled_model_onnx_gpu([image])
    end = time.perf_counter()
    time_onnx_gpu = end - start
    print(
        f"ONNX model in OpenVINO/GPU: {time_onnx_gpu/num_images:.3f} "
        f"seconds per image, FPS: {num_images/time_onnx_gpu:.2f}"
    )

    compiled_model_ir_gpu = ie.compile_model(model=model_ir, device_name="GPU")
    start = time.perf_counter()
    for _ in range(num_images):
        compiled_model_ir_gpu([image])
    end = time.perf_counter()
    time_ir_gpu = end - start
    print(
        f"IR model in OpenVINO/GPU: {time_ir_gpu/num_images:.3f} "
        f"seconds per image, FPS: {num_images/time_ir_gpu:.2f}"
    )

and the output

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
Cell In[3], line 10
      8 from detectron2.modeling import build_model
      9 from detectron2.checkpoint import DetectionCheckpointer
---> 10 from detectron2.export.caffe2_export import export_onnx_model
     11 import onnx
     12 import onnxruntime as rt

File [~\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.10_qbz5n2kfra8p0\LocalCache\local-packages\Python310\site-packages\detectron2\export\caffe2_export.py:9](https://file+.vscode-resource.vscode-cdn.net/c%3A/Users/Lazhar%20Bouacha/OneDrive%20-%20CHAMP%27S/PhD/~/AppData/Local/Packages/PythonSoftwareFoundation.Python.3.10_qbz5n2kfra8p0/LocalCache/local-packages/Python310/site-packages/detectron2/export/caffe2_export.py:9)
      7 from typing import List
      8 import onnx
----> 9 import onnx.optimizer
     10 import torch
     11 from caffe2.proto import caffe2_pb2

ModuleNotFoundError: No module named 'onnx.optimizer'

[jcwchen]

Thanks for the details. I think it seems like an issue in detectron2: https://github.com/facebookresearch/detectron2/blob/3ed66980529aadf662f469bca744221cab762e83/detectron2/export/caffe2_export.py#L9. onnx.optimizer has moved out since ONNX 1.9.0. You can still find onnx.optimizer with ONNX < 1.9.0, but going forward you should use onnxoptimizer from https://github.com/onnx/optimizer. It should be fixed from detectron2 side if you really want to use the latest ONNX.