exporter.py 4.91 KB
Newer Older
facebook-github-bot's avatar
facebook-github-bot committed
1
2
3
4
5
6
7
8
9
10
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved

"""
Binary to convert pytorch detectron2go model to a predictor, which contains model(s) in
deployable format (such as torchscript, caffe2, ...)
"""

import copy
import logging
11
import sys
12
13
from dataclasses import dataclass
from typing import Any, Dict, List, Type, Union
facebook-github-bot's avatar
facebook-github-bot committed
14
15

import mobile_cv.lut.lib.pt.flops_utils as flops_utils
16
from d2go.config import CfgNode, temp_defrost
17
from d2go.distributed import set_shared_context
18
from d2go.export.exporter import convert_and_export_predictor
19
from d2go.runner import BaseRunner
20
21
22
23
24
from d2go.setup import (
    basic_argument_parser,
    post_mortem_if_fail_for_main,
    prepare_for_launch,
    setup_after_launch,
25
    setup_before_launch,
26
    setup_root_logger,
27
)
facebook-github-bot's avatar
facebook-github-bot committed
28
29
30
31
32


logger = logging.getLogger("d2go.tools.export")


33
34
35
36
37
38
@dataclass
class ExporterOutput:
    predictor_paths: Dict[str, str]
    accuracy_comparison: Dict[str, Any]


facebook-github-bot's avatar
facebook-github-bot committed
39
def main(
40
41
42
    cfg: CfgNode,
    output_dir: str,
    runner_class: Union[str, Type[BaseRunner]],
facebook-github-bot's avatar
facebook-github-bot committed
43
    # binary specific optional arguments
44
    predictor_types: List[str],
45
    device: str = "cpu",
facebook-github-bot's avatar
facebook-github-bot committed
46
47
    compare_accuracy: bool = False,
    skip_if_fail: bool = False,
48
    skip_model_weights: bool = False,
49
) -> ExporterOutput:
50
51
52
53
54
55
56
    if compare_accuracy:
        raise NotImplementedError(
            "compare_accuracy functionality isn't currently supported."
        )
        # NOTE: dict for metrics of all exported models (and original pytorch model)
        # ret["accuracy_comparison"] = accuracy_comparison

facebook-github-bot's avatar
facebook-github-bot committed
57
    cfg = copy.deepcopy(cfg)
58
59
60
61
    with temp_defrost(cfg):
        if skip_model_weights:
            cfg.merge_from_list(["MODEL.WEIGHTS", ""])

62
    runner = setup_after_launch(cfg, output_dir, runner_class)
facebook-github-bot's avatar
facebook-github-bot committed
63
64

    with temp_defrost(cfg):
65
        cfg.merge_from_list(["MODEL.DEVICE", device])
facebook-github-bot's avatar
facebook-github-bot committed
66

67
    model = runner.build_model(cfg, eval_only=True)
facebook-github-bot's avatar
facebook-github-bot committed
68
69
70
    # NOTE: train dataset is used to avoid leakage since the data might be used for
    # running calibration for quantization. test_loader is used to make sure it follows
    # the inference behaviour (augmentation will not be applied).
71
    datasets = list(cfg.DATASETS.TRAIN)
RangiLyu's avatar
RangiLyu committed
72
    data_loader = runner.build_detection_test_loader(cfg, datasets)
facebook-github-bot's avatar
facebook-github-bot committed
73
74
75
76
77
78

    logger.info("Running the pytorch model and print FLOPS ...")
    first_batch = next(iter(data_loader))
    input_args = (first_batch,)
    flops_utils.print_model_flops(model, input_args)

79
    predictor_paths: Dict[str, str] = {}
facebook-github-bot's avatar
facebook-github-bot committed
80
81
82
83
84
    for typ in predictor_types:
        # convert_and_export_predictor might alter the model, copy before calling it
        pytorch_model = copy.deepcopy(model)
        try:
            predictor_path = convert_and_export_predictor(
85
86
87
88
89
                cfg,
                pytorch_model,
                typ,
                output_dir,
                data_loader,
facebook-github-bot's avatar
facebook-github-bot committed
90
91
92
93
            )
            logger.info(f"Predictor type {typ} has been exported to {predictor_path}")
            predictor_paths[typ] = predictor_path
        except Exception as e:
94
            logger.exception(f"Export {typ} predictor failed: {e}")
facebook-github-bot's avatar
facebook-github-bot committed
95
96
97
            if not skip_if_fail:
                raise e

98
99
100
101
    return ExporterOutput(
        predictor_paths=predictor_paths,
        accuracy_comparison={},
    )
facebook-github-bot's avatar
facebook-github-bot committed
102
103
104


def run_with_cmdline_args(args):
105
    cfg, output_dir, runner_name = prepare_for_launch(args)
106
107
108
109
    shared_context = setup_before_launch(cfg, output_dir, runner_name)
    if shared_context is not None:
        set_shared_context(shared_context)

110
111
    main_func = main if args.disable_post_mortem else post_mortem_if_fail_for_main(main)
    return main_func(
facebook-github-bot's avatar
facebook-github-bot committed
112
113
        cfg,
        output_dir,
114
        runner_name,
facebook-github-bot's avatar
facebook-github-bot committed
115
116
        # binary specific optional arguments
        predictor_types=args.predictor_types,
117
        device=args.device,
facebook-github-bot's avatar
facebook-github-bot committed
118
119
        compare_accuracy=args.compare_accuracy,
        skip_if_fail=args.skip_if_fail,
120
        skip_model_weights=args.skip_model_weights,
facebook-github-bot's avatar
facebook-github-bot committed
121
122
123
124
125
126
127
128
129
130
131
    )


def get_parser():
    parser = basic_argument_parser(distributed=False)
    parser.add_argument(
        "--predictor-types",
        type=str,
        nargs="+",
        help="List of strings specify the types of predictors to export",
    )
132
133
134
    parser.add_argument(
        "--device", default="cpu", help="the device to export the model on"
    )
facebook-github-bot's avatar
facebook-github-bot committed
135
136
137
138
    parser.add_argument(
        "--compare-accuracy",
        action="store_true",
        help="If true, all exported models and the original pytorch model will be"
Alexander Pivovarov's avatar
Alexander Pivovarov committed
139
        " evaluated on cfg.DATASETS.TEST",
facebook-github-bot's avatar
facebook-github-bot committed
140
141
142
143
144
145
146
147
    )
    parser.add_argument(
        "--skip-if-fail",
        action="store_true",
        default=False,
        help="If set, suppress the exception for failed exporting and continue to"
        " export the next type of model",
    )
148
149
    parser.add_argument("--skip-model-weights", action="store_true")

facebook-github-bot's avatar
facebook-github-bot committed
150
151
    return parser

152

Tsahi Glik's avatar
Tsahi Glik committed
153
154
def cli(args=None):
    args = sys.argv[1:] if args is None else args
155
    run_with_cmdline_args(get_parser().parse_args(args))
facebook-github-bot's avatar
facebook-github-bot committed
156

157

facebook-github-bot's avatar
facebook-github-bot committed
158
if __name__ == "__main__":
159
    setup_root_logger()
Tsahi Glik's avatar
Tsahi Glik committed
160
    cli()