export API & tests

Reviewed By: wat3rBro

Differential Revision: D19255042

fbshipit-source-id: 6b12bba55c1a86fa65868d235aead7f393133c19

detectron2/export/__init__.py

@@ -0,0 +1,3 @@
+# -*- coding: utf-8 -*-
+
+from .api import *

detectron2/export/api.py

@@ -0,0 +1,152 @@
+import logging
+import os
+import torch
+from caffe2.proto import caffe2_pb2
+from torch import nn
+
+from detectron2.config import CfgNode as CN
+
+from .caffe2_export import export_caffe2_detection_model, run_and_save_graph
+from .caffe2_inference import ProtobufDetectionModel
+from .caffe2_modeling import META_ARCH_CAFFE2_EXPORT_TYPE_MAP, convert_batched_inputs_to_c2_format
+from .shared import get_pb_arg_vali, save_graph
+
+__all__ = ["add_export_config", "export_caffe2_model", "Caffe2Model"]
+
+
+def add_export_config(cfg):
+    """
+    Args:
+        cfg (CfgNode): a detectron2 config
+
+    Returns:
+        CfgNode: an updated config with new options that :func:`export_caffe2_model` will need.
+    """
+    is_frozen = cfg.is_frozen()
+    cfg.defrost()
+    cfg.EXPORT_CAFFE2 = CN()
+    cfg.EXPORT_CAFFE2.USE_HEATMAP_MAX_KEYPOINT = False
+    if is_frozen:
+        cfg.freeze()
+    return cfg
+
+
+def export_caffe2_model(cfg, model, inputs):
+    """
+    Export a detectron2 model to caffe2 format.
+
+    Args:
+        cfg (CfgNode): a detectron2 config, with extra export-related options
+            added by :func:`add_export_config`.
+        model (nn.Module): a model built by
+            :func:`detectron2.modeling.build_model`.
+            It will be modified by this function.
+        inputs: sample inputs that the given model takes for inference.
+            Will be used to trace the model.
+
+    Returns:
+        Caffe2Model
+    """
+    assert isinstance(cfg, CN), cfg
+    C2MetaArch = META_ARCH_CAFFE2_EXPORT_TYPE_MAP[cfg.MODEL.META_ARCHITECTURE]
+    c2_compatible_model = C2MetaArch(cfg, model)
+    c2_format_input = c2_compatible_model.get_tensors_input(inputs)
+    predict_net, init_net = export_caffe2_detection_model(c2_compatible_model, c2_format_input)
+    return Caffe2Model(predict_net, init_net)
+
+
+class Caffe2Model(nn.Module):
+    def __init__(self, predict_net, init_net):
+        super().__init__()
+        self.eval()  # always in eval mode
+        self._predict_net = predict_net
+        self._init_net = init_net
+        self._predictor = None
+
+    @property
+    def predict_net(self):
+        """
+        Returns:
+            core.Net: the underlying caffe2 predict net
+        """
+        return self._predict_net
+
+    @property
+    def init_net(self):
+        """
+        Returns:
+            core.Net: the underlying caffe2 init net
+        """
+        return self._init_net
+
+    __init__.__HIDE_SPHINX_DOC__ = True
+
+    def save_protobuf(self, output_dir):
+        """
+        Save the model as caffe2's protobuf format.
+
+        Args:
+            output_dir (str): the output directory to save protobuf files.
+        """
+        logger = logging.getLogger(__name__)
+        logger.info("Saving model to {} ...".format(output_dir))
+        os.makedirs(output_dir, exist_ok=True)
+
+        with open(os.path.join(output_dir, "model.pb"), "wb") as f:
+            f.write(self._predict_net.SerializeToString())
+        with open(os.path.join(output_dir, "model.pbtxt"), "w") as f:
+            f.write(str(self._predict_net))
+        with open(os.path.join(output_dir, "model_init.pb"), "wb") as f:
+            f.write(self._init_net.SerializeToString())
+
+    def save_graph(self, output_file, inputs=None):
+        """
+        Save the graph as SVG format.
+
+        Args:
+            output_file (str): a SVG file
+            inputs: optional inputs given to the model.
+                If given, the inputs will be used to run the graph to record
+                shape of every tensor. The shape information will be
+                saved together with the graph.
+        """
+        if inputs is None:
+            save_graph(self._predict_net, output_file, op_only=False)
+        else:
+            size_divisibility = get_pb_arg_vali(self._predict_net, "size_divisibility", 0)
+            inputs = convert_batched_inputs_to_c2_format(
+                inputs, size_divisibility, torch.device("cpu")
+            )
+            inputs = [x.numpy() for x in inputs]
+            run_and_save_graph(self._predict_net, self._init_net, inputs, output_file)
+
+    @staticmethod
+    def load_protobuf(dir):
+        """
+        Args:
+            dir (str): a directory used to save Caffe2Model with
+                :meth:`save_protobuf`.
+                The files "model.pb" and "model_init.pb" are needed.
+
+        Returns:
+            Caffe2Model: the caffe2 model loaded from this directory.
+        """
+        predict_net = caffe2_pb2.NetDef()
+        with open(os.path.join(dir, "model.pb"), "rb") as f:
+            predict_net.ParseFromString(f.read())
+
+        init_net = caffe2_pb2.NetDef()
+        with open(os.path.join(dir, "model_init.pb"), "rb") as f:
+            init_net.ParseFromString(f.read())
+
+        return Caffe2Model(predict_net, init_net)
+
+    def __call__(self, inputs):
+        """
+        An interface that wraps around a caffe2 model and mimics detectron2's models'
+        input & output format. This is used to compare the caffe2 model
+        with its original torch model.
+        """
+        if self._predictor is None:
+            self._predictor = ProtobufDetectionModel(self._predict_net, self._init_net)
+        return self._predictor(inputs)

detectron2/export/caffe2_export.py

@@ -76,7 +76,7 @@ def export_caffe2_detection_model(model: torch.nn.Module, tensor_inputs: List[to
     assert hasattr(model, "encode_additional_info")
 
     # Export via ONNX
-    logger.info("Exporting {} model via ONNX ...".format(type(model)))
+    logger.info("Exporting a {} model via ONNX ...".format(type(model).__name__))
     predict_net, init_net = _export_via_onnx(model, (tensor_inputs,))
     logger.info("ONNX export Done.")
 

detectron2/export/caffe2_inference.py

@@ -139,8 +139,8 @@ def assemble(self, batched_inputs, c2_inputs, c2_results):
         for sem_seg_result, detector_result, input_per_image, image_size in zip(
             sem_seg_results, detector_results, batched_inputs, image_sizes
         ):
-            height = input_per_image.get("height")
-            width = input_per_image.get("width")
+            height = input_per_image.get("height", image_size[0])
+            width = input_per_image.get("width", image_size[1])
             sem_seg_r = sem_seg_postprocess(sem_seg_result, image_size, height, width)
             detector_r = detector_postprocess(detector_result, height, width)
 

detectron2/export/caffe2_modeling.py

@@ -2,14 +2,13 @@
 
 import contextlib
 import io
+import mock
 import struct
 import torch
 
 from detectron2.modeling import meta_arch
 from detectron2.structures import ImageList
 
-import mock
-
 from .c10 import Caffe2Compatible
 from .patcher import ROIHeadsPatcher, patch_generalized_rcnn
 from .shared import alias, check_set_pb_arg, mock_torch_nn_functional_interpolate

detectron2/export/patcher.py

@@ -1,15 +1,14 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
 
 import contextlib
+import mock
 import torch
 
 from detectron2.modeling import poolers
 from detectron2.modeling.proposal_generator import rpn
 from detectron2.modeling.roi_heads import roi_heads
 from detectron2.modeling.roi_heads.fast_rcnn import FastRCNNOutputs
 
-import mock
-
 from .c10 import (
     Caffe2Compatible,
     Caffe2FastRCNNOutputsInference,

detectron2/export/shared.py

@@ -5,6 +5,7 @@
 import copy
 import functools
 import logging
+import mock
 import numpy as np
 import os
 from typing import Any, Dict, List, Optional, Tuple, Union
@@ -15,8 +16,6 @@
 from caffe2.python import core, net_drawer, workspace
 from torch.nn.functional import interpolate as interp
 
-import mock
-
 logger = logging.getLogger(__name__)
 
 

detectron2/modeling/meta_arch/panoptic_fpn.py

@@ -109,8 +109,8 @@ def forward(self, batched_inputs):
         for sem_seg_result, detector_result, input_per_image, image_size in zip(
             sem_seg_results, detector_results, batched_inputs, images.image_sizes
         ):
-            height = input_per_image.get("height")
-            width = input_per_image.get("width")
+            height = input_per_image.get("height", image_size[0])
+            width = input_per_image.get("width", image_size[1])
             sem_seg_r = sem_seg_postprocess(sem_seg_result, image_size, height, width)
             detector_r = detector_postprocess(detector_result, height, width)
 

setup.cfg

@@ -2,7 +2,7 @@
 line_length=100
 multi_line_output=3
 include_trailing_comma=True
-known_standard_library=numpy,setuptools
+known_standard_library=numpy,setuptools,mock
 skip=datasets,docs
 skip_glob=*/__init__.py
 known_myself=detectron2

tools/caffe2_converter.py

@@ -0,0 +1,60 @@
+import argparse
+import os
+
+from detectron2.checkpoint import DetectionCheckpointer
+from detectron2.config import get_cfg
+from detectron2.data import build_detection_test_loader
+from detectron2.evaluation import COCOEvaluator, inference_on_dataset, print_csv_format
+from detectron2.export import add_export_config, export_caffe2_model
+from detectron2.modeling import build_model
+from detectron2.utils.logger import setup_logger
+
+
+def setup_cfg(args):
+    cfg = get_cfg()
+    cfg = add_export_config(cfg)
+    cfg.merge_from_file(args.config_file)
+    cfg.merge_from_list(args.opts)
+    cfg.freeze()
+    return cfg
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Convert a model to Caffe2")
+    parser.add_argument("--config-file", default="", metavar="FILE", help="path to config file")
+    parser.add_argument("--run-eval", action="store_true")
+    parser.add_argument("--output", help="output directory for the converted caffe2 model")
+    parser.add_argument(
+        "opts",
+        help="Modify config options using the command-line",
+        default=None,
+        nargs=argparse.REMAINDER,
+    )
+    args = parser.parse_args()
+    logger = setup_logger()
+    logger.info("Command line arguments: " + str(args))
+
+    cfg = setup_cfg(args)
+
+    # create a torch model
+    torch_model = build_model(cfg)
+    DetectionCheckpointer(torch_model).resume_or_load(cfg.MODEL.WEIGHTS)
+
+    # get a sample data
+    data_loader = build_detection_test_loader(cfg, cfg.DATASETS.TEST[0])
+    first_batch = next(iter(data_loader))
+
+    # convert and save caffe2 model
+    caffe2_model = export_caffe2_model(cfg, torch_model, first_batch)
+    caffe2_model.save_protobuf(args.output)
+    # draw the caffe2 graph
+    caffe2_model.save_graph(os.path.join(args.output, "model_def.svg"), inputs=first_batch)
+
+    # run evaluation with the converted model
+    if args.run_eval:
+        dataset = cfg.DATASETS.TEST[0]
+        data_loader = build_detection_test_loader(cfg, dataset)
+        # NOTE: hard-coded evaluator. change to the evaluator for your dataset
+        evaluator = COCOEvaluator(dataset, cfg, True, args.output)
+        metrics = inference_on_dataset(caffe2_model, data_loader, evaluator)
+        print_csv_format(metrics)