exporter.py 4.59 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
)
facebook-github-bot's avatar
facebook-github-bot committed
27
28
29
30
31


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


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


facebook-github-bot's avatar
facebook-github-bot committed
38
def main(
39
40
41
    cfg: CfgNode,
    output_dir: str,
    runner_class: Union[str, Type[BaseRunner]],
facebook-github-bot's avatar
facebook-github-bot committed
42
    # binary specific optional arguments
43
    predictor_types: List[str],
44
    device: str = "cpu",
facebook-github-bot's avatar
facebook-github-bot committed
45
46
    compare_accuracy: bool = False,
    skip_if_fail: bool = False,
47
) -> ExporterOutput:
48
49
50
51
52
53
54
    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
55
    cfg = copy.deepcopy(cfg)
56
    runner = setup_after_launch(cfg, output_dir, runner_class)
facebook-github-bot's avatar
facebook-github-bot committed
57
58

    with temp_defrost(cfg):
59
        cfg.merge_from_list(["MODEL.DEVICE", device])
facebook-github-bot's avatar
facebook-github-bot committed
60
61
62
63
64
    model = runner.build_model(cfg, eval_only=True)

    # 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).
65
    datasets = list(cfg.DATASETS.TRAIN)
RangiLyu's avatar
RangiLyu committed
66
    data_loader = runner.build_detection_test_loader(cfg, datasets)
facebook-github-bot's avatar
facebook-github-bot committed
67
68
69
70
71
72

    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)

73
    predictor_paths: Dict[str, str] = {}
facebook-github-bot's avatar
facebook-github-bot committed
74
75
76
77
78
    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(
79
80
81
82
83
                cfg,
                pytorch_model,
                typ,
                output_dir,
                data_loader,
facebook-github-bot's avatar
facebook-github-bot committed
84
85
86
87
            )
            logger.info(f"Predictor type {typ} has been exported to {predictor_path}")
            predictor_paths[typ] = predictor_path
        except Exception as e:
88
            logger.exception(f"Export {typ} predictor failed: {e}")
facebook-github-bot's avatar
facebook-github-bot committed
89
90
91
            if not skip_if_fail:
                raise e

92
93
94
95
    return ExporterOutput(
        predictor_paths=predictor_paths,
        accuracy_comparison={},
    )
facebook-github-bot's avatar
facebook-github-bot committed
96
97
98


def run_with_cmdline_args(args):
99
    cfg, output_dir, runner_name = prepare_for_launch(args)
100
101
102
103
    shared_context = setup_before_launch(cfg, output_dir, runner_name)
    if shared_context is not None:
        set_shared_context(shared_context)

104
105
    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
106
107
        cfg,
        output_dir,
108
        runner_name,
facebook-github-bot's avatar
facebook-github-bot committed
109
110
        # binary specific optional arguments
        predictor_types=args.predictor_types,
111
        device=args.device,
facebook-github-bot's avatar
facebook-github-bot committed
112
113
114
115
116
117
118
119
120
121
122
123
124
        compare_accuracy=args.compare_accuracy,
        skip_if_fail=args.skip_if_fail,
    )


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",
    )
125
126
127
    parser.add_argument(
        "--device", default="cpu", help="the device to export the model on"
    )
facebook-github-bot's avatar
facebook-github-bot committed
128
129
130
131
    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
132
        " evaluated on cfg.DATASETS.TEST",
facebook-github-bot's avatar
facebook-github-bot committed
133
134
135
136
137
138
139
140
141
142
    )
    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",
    )
    return parser

143

Tsahi Glik's avatar
Tsahi Glik committed
144
145
def cli(args=None):
    args = sys.argv[1:] if args is None else args
146
    run_with_cmdline_args(get_parser().parse_args(args))
facebook-github-bot's avatar
facebook-github-bot committed
147

148

facebook-github-bot's avatar
facebook-github-bot committed
149
if __name__ == "__main__":
Tsahi Glik's avatar
Tsahi Glik committed
150
    cli()