Function main

examples/apple/coreml/scripts/debugger_cli.py:131–177

()

def main() -> None:
    parser = args_parser()
    args = parser.parse_args()

    if args.model_name not in MODEL_NAME_TO_MODEL:
        raise RuntimeError(
            f"Model {args.model_name} is not a valid name. "
            f"Available models are {list(MODEL_NAME_TO_MODEL.keys())}."
        )

    valid_compute_units = [compute_unit.name.lower() for compute_unit in ct.ComputeUnit]
    if args.compute_unit not in valid_compute_units:
        raise RuntimeError(
            f"{args.compute_unit} is invalid. "
            f"Valid compute units are {valid_compute_units}."
        )

    build_devtools_runner_including_coreml(
        root_dir_path=get_root_dir_path(), conda_env_name=args.conda_environment_name
    )

    model, example_inputs, _, _ = EagerModelFactory.create_model(
        *MODEL_NAME_TO_MODEL[args.model_name]
    )

    model.eval()
    edge_program = module_to_edge(
        module=model,
        example_inputs=example_inputs,
    )

    coreml_compile_specs = get_compile_specs_from_args(args)

    with tempfile.TemporaryDirectory() as temp_dir_name:
        working_dir_path = Path(temp_dir_name) / "debugger"
        working_dir_path.mkdir(parents=True, exist_ok=True)
        comparison_result = compare_intermediate_tensors(
            edge_program=edge_program,
            example_inputs=example_inputs,
            coreml_compile_specs=coreml_compile_specs,
            model_name=args.model_name,
            working_dir_path=working_dir_path,
        )

        print(
            tabulate(comparison_result.to_dataframe(), headers="keys", tablefmt="grid")
        )


if __name__ == "__main__":