lufficc / SSD

High quality, fast, modular reference implementation of SSD in PyTorch
MIT License
1.52k stars 384 forks source link

how to export onnx and openvino? #228

Closed 1485995573 closed 5 months ago

1485995573 commented 5 months ago
import openvino as ov
import torch
from PIL import Image
from ssd.config import cfg
from ssd.modeling.detector import build_detection_model
import torchvision.transforms as transforms
import numpy as np
import time

config_file = 'configs/mobilenet_v2_ssd320_voc0712.yaml'
ckpt_path = 'mobilenet_v2_ssd320_voc0712_v2.pth'
image_path = '/home/f/mycode/datasets/coco/images/val2017/000000000285.jpg'

cfg.merge_from_file(config_file)
cfg.freeze()
model = build_detection_model(cfg)
model.eval()

# 定义图像转换
transform = transforms.Compose([
    transforms.Resize((320, 320)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 定义输入尺寸
input_shape = (1, 3, 320, 320)  # 根据模型的输入尺寸设定
img = torch.randn(input_shape)
# img = transform(img).unsqueeze(0)
# # 将输入数据移动到所选设备上
# img_batch = img.repeat(1, 1, 1, 1)
# # 创建一个虚拟输入张量

# 导出模型为 ONNX 格式
onnx_path = 'mobilenet_v2_ssd320_voc0712_1.onnx'
torch.onnx.export(model, img, onnx_path, verbose=True, export_params=True, opset_version=11)

# 转换 ONNX 模型为 OpenVINO 格式
ov_model = ov.convert_model(onnx_path)

# 保存 OpenVINO 模型
ir_path = 'mobilenet_v2_ssd320_voc0712_openvino_1/mobilenet_v2_ssd320_voc0712_1.xml'
ov.save_model(ov_model, ir_path)

print("OpenVINO IR model saved to:", ir_path)

my issue here👇 RuntimeError: Only tuples, lists and Variables are supported as JIT inputs/outputs. Dictionaries and strings are also accepted, but their usage is not recommended. Here, received an input of unsupported type: Container