def export_model_to_onnx()

in optimum/furiosa/utils.py [0:0]

• 12 lines of code
• 1 McCabe index (conditional complexity)

def export_model_to_onnx(model_id, save_dir, input_shape_dict, output_shape_dict, file_name="model.onnx"):
    task = "image-classification"
    main_export(model_id, save_dir, task=task)

    import onnx
    from onnx import shape_inference
    from onnx.tools import update_model_dims

    save_dir_path = Path(save_dir) / "model.onnx"
    model = onnx.load(save_dir_path)
    updated_model = update_model_dims.update_inputs_outputs_dims(model, input_shape_dict, output_shape_dict)
    inferred_model = shape_inference.infer_shapes(updated_model)

    static_model_path = Path(save_dir_path).parent / file_name
    onnx.save(inferred_model, static_model_path)