Skip to content

I cannot understand the difference between convert onnx by torch.onnx.export and torch2onnx(mmdeploy) #3856

Description

@Matthew-DgNg

*** I cannot understand the difference between convert onnx by torch.onnx.export and torch2onnx(mmdeploy) (mmdeploy have better result). The inference results are differents, and each model cannot run in other case framework( eg: model convert from pytorch.onnx.export cannot run on mmdeploy inference and oppostie way is also the same)

-- pytorch:

def convertModel(config_file: str, checkpoint_file: str, names, onnx_path: str, ir_path: str, IMAGE_WIDTH: any = None, IMAGE_HEIGHT: any = None):
# Load config
cfg = Config.fromfile(config_file)
cfg.model.pretrained = None

# Build model
model = build_segmentor(cfg.model)
model.cfg = cfg
model.eval()

# Load state dict thủ công (tránh lỗi pickle)
ckpt = torch.load(checkpoint_file, map_location='cpu', weights_only=False)
model.load_state_dict(ckpt['state_dict'])  # hoặc ckpt nếu không có 'state_dict'
print("Loaded model")
metadata = {
        "names": names,
        "type": "MmSegmentation"
    }
with warnings.catch_warnings():
    warnings.filterwarnings("ignore")
    dummy_input = torch.randn(1, 3, IMAGE_HEIGHT, IMAGE_WIDTH)
    os.makedirs(os.path.dirname(onnx_path), exist_ok=True)
    torch.onnx.export(
        model,
        dummy_input,
        onnx_path,
        input_names=['input'],
        output_names=['output'],
        opset_version=14,
        do_constant_folding=True,
        verbose=True
    )
model_onnx = onnx.load(onnx_path)  # load onnx model

for k, v in metadata.items():
    meta = model_onnx.metadata_props.add()
    meta.key, meta.value = k, str(v)

def normalize(image: np.ndarray) -> np.ndarray:
image = image.astype(np.float16)
mean = (103.53, 116.28, 123.675)
std = (57.375, 57.12, 58.395)
image -= mean
image /= std
return image

def inference(config_file, checkpoint_file, onnx_path, ir_path, image_filename: str, label_map, device: str = "CPU", IMAGE_WIDTH: any = None, IMAGE_HEIGHT: any = None):

# image = cv2.cvtColor(cv2.imread(str(image_filename)), cv2.COLOR_BGR2RGB)

image = cv2.imread(str(image_filename))
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
resized_image = cv2.resize(image, (IMAGE_WIDTH, IMAGE_HEIGHT))
normalized_image = normalize(resized_image)
print("Mean after normalize:", normalized_image.mean(axis=(0, 1)))
print("Std after normalize:", normalized_image.std(axis=(0, 1)))


# Convert the resized images to network input shape.
normalized_input_image = np.expand_dims(np.transpose(normalized_image, (2, 0, 1)), 0)
print("Shape:", normalized_input_image.shape)
print("Dtype:", normalized_input_image.dtype)

# Instantiate OpenVINO Core
core = ov.Core()

# Read model to OpenVINO Runtime
model_onnx = core.read_model(model=onnx_path)

# Load model on device
compiled_model_onnx = core.compile_model(model=model_onnx, device_name=device)

# Run inference on the input image
res_onnx = compiled_model_onnx([normalized_input_image])[0]
Labels = SegmentationMap(label_map)

# Convert the network result to a segmentation map and display the result.
result_mask_onnx = np.squeeze(np.argmax(res_onnx, axis=1)).astype(np.uint8)
viz_result_image(
    image,
    segmentation_map_to_image(result_mask_onnx, Labels.get_colormap()),
    resize=True,
)
Image

-- mmdeploy:

ckpt = torch.load(model_checkpoint, map_location='cuda', weights_only=False)
torch.save(ckpt['state_dict'], 'work_dir/onnx/fullFoam1024/iter_17800.pth')

torch2onnx(img, work_dir, save_file, deploy_cfg, model_cfg,
'work_dir/onnx/fullFoam1024/iter_17800.pth', device)

onnx_model = onnx.load('work_dir\onnx\fullFoam1024\model.onnx')

Add metadata properties

metadata = {
"names": "{0: 'none', 1: 'foam'}",
"type": "MmSegmentation"
}

for k, v in metadata.items():
meta = onnx_model.metadata_props.add()
meta.key, meta.value = k, str(v)

onnx.save(onnx_model, 'work_dir\onnx\fullFoam1024\model.onnx')
""" Inference onnx segmentation model"""

read deploy_cfg and model_cfg

deploy_cfg, model_cfg = load_config(deploy_cfg, model_cfg)

build task and backend model

task_processor = build_task_processor(model_cfg, deploy_cfg, device='cpu')
model = task_processor.build_backend_model(onnx_model)

process input image

input_shape = get_input_shape(deploy_cfg)
model_inputs, _ = task_processor.create_input(img, input_shape)

do model inference

with torch.no_grad():
start_time = time.perf_counter()
result = model.test_step(model_inputs)
elapsed_time = (time.perf_counter() - start_time) * 1000
print(f"Inference time: {elapsed_time:.2f} ms")
output = result[0].pred_sem_seg.data.cpu().numpy()
print("Unique labels in result:", np.unique(output)) # Debug

visualize results

task_processor.visualize(
image=img,
model=model,
result=result[0],
window_name='visualize',
output_file='work_dir/onnx/fullFoam1024/output_segmentation1.png')

Image

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions