Arulkumar03/Wheat_HEAD_Detection_Counting_ComputerVision_Model

1# Copyright (c) Facebook, Inc. and its affiliates.2import copy3import logging4import os5import torch6from caffe2.proto import caffe2_pb27from torch import nn8 9from detectron2.config import CfgNode10from detectron2.utils.file_io import PathManager11 12from .caffe2_inference import ProtobufDetectionModel13from .caffe2_modeling import META_ARCH_CAFFE2_EXPORT_TYPE_MAP, convert_batched_inputs_to_c2_format14from .shared import get_pb_arg_vali, get_pb_arg_vals, save_graph15 16__all__ = [17    "Caffe2Model",18    "Caffe2Tracer",19]20 21 22class Caffe2Tracer:23    """24    Make a detectron2 model traceable with Caffe2 operators.25    This class creates a traceable version of a detectron2 model which:26 27    1. Rewrite parts of the model using ops in Caffe2. Note that some ops do28       not have GPU implementation in Caffe2.29    2. Remove post-processing and only produce raw layer outputs30 31    After making a traceable model, the class provide methods to export such a32    model to different deployment formats.33    Exported graph produced by this class take two input tensors:34 35    1. (1, C, H, W) float "data" which is an image (usually in [0, 255]).36       (H, W) often has to be padded to multiple of 32 (depend on the model37       architecture).38    2. 1x3 float "im_info", each row of which is (height, width, 1.0).39       Height and width are true image shapes before padding.40 41    The class currently only supports models using builtin meta architectures.42    Batch inference is not supported, and contributions are welcome.43    """44 45    def __init__(self, cfg: CfgNode, model: nn.Module, inputs):46        """47        Args:48            cfg (CfgNode): a detectron2 config used to construct caffe2-compatible model.49            model (nn.Module): An original pytorch model. Must be among a few official models50                in detectron2 that can be converted to become caffe2-compatible automatically.51                Weights have to be already loaded to this model.52            inputs: sample inputs that the given model takes for inference.53                Will be used to trace the model. For most models, random inputs with54                no detected objects will not work as they lead to wrong traces.55        """56        assert isinstance(cfg, CfgNode), cfg57        assert isinstance(model, torch.nn.Module), type(model)58 59        # TODO make it support custom models, by passing in c2 model directly60        C2MetaArch = META_ARCH_CAFFE2_EXPORT_TYPE_MAP[cfg.MODEL.META_ARCHITECTURE]61        self.traceable_model = C2MetaArch(cfg, copy.deepcopy(model))62        self.inputs = inputs63        self.traceable_inputs = self.traceable_model.get_caffe2_inputs(inputs)64 65    def export_caffe2(self):66        """67        Export the model to Caffe2's protobuf format.68        The returned object can be saved with its :meth:`.save_protobuf()` method.69        The result can be loaded and executed using Caffe2 runtime.70 71        Returns:72            :class:`Caffe2Model`73        """74        from .caffe2_export import export_caffe2_detection_model75 76        predict_net, init_net = export_caffe2_detection_model(77            self.traceable_model, self.traceable_inputs78        )79        return Caffe2Model(predict_net, init_net)80 81    def export_onnx(self):82        """83        Export the model to ONNX format.84        Note that the exported model contains custom ops only available in caffe2, therefore it85        cannot be directly executed by other runtime (such as onnxruntime or TensorRT).86        Post-processing or transformation passes may be applied on the model to accommodate87        different runtimes, but we currently do not provide support for them.88 89        Returns:90            onnx.ModelProto: an onnx model.91        """92        from .caffe2_export import export_onnx_model as export_onnx_model_impl93 94        return export_onnx_model_impl(self.traceable_model, (self.traceable_inputs,))95 96    def export_torchscript(self):97        """98        Export the model to a ``torch.jit.TracedModule`` by tracing.99        The returned object can be saved to a file by ``.save()``.100 101        Returns:102            torch.jit.TracedModule: a torch TracedModule103        """104        logger = logging.getLogger(__name__)105        logger.info("Tracing the model with torch.jit.trace ...")106        with torch.no_grad():107            return torch.jit.trace(self.traceable_model, (self.traceable_inputs,))108 109 110class Caffe2Model(nn.Module):111    """112    A wrapper around the traced model in Caffe2's protobuf format.113    The exported graph has different inputs/outputs from the original Pytorch114    model, as explained in :class:`Caffe2Tracer`. This class wraps around the115    exported graph to simulate the same interface as the original Pytorch model.116    It also provides functions to save/load models in Caffe2's format.'117 118    Examples:119    ::120        c2_model = Caffe2Tracer(cfg, torch_model, inputs).export_caffe2()121        inputs = [{"image": img_tensor_CHW}]122        outputs = c2_model(inputs)123        orig_outputs = torch_model(inputs)124    """125 126    def __init__(self, predict_net, init_net):127        super().__init__()128        self.eval()  # always in eval mode129        self._predict_net = predict_net130        self._init_net = init_net131        self._predictor = None132 133    __init__.__HIDE_SPHINX_DOC__ = True134 135    @property136    def predict_net(self):137        """138        caffe2.core.Net: the underlying caffe2 predict net139        """140        return self._predict_net141 142    @property143    def init_net(self):144        """145        caffe2.core.Net: the underlying caffe2 init net146        """147        return self._init_net148 149    def save_protobuf(self, output_dir):150        """151        Save the model as caffe2's protobuf format.152        It saves the following files:153 154            * "model.pb": definition of the graph. Can be visualized with155              tools like `netron <https://github.com/lutzroeder/netron>`_.156            * "model_init.pb": model parameters157            * "model.pbtxt": human-readable definition of the graph. Not158              needed for deployment.159 160        Args:161            output_dir (str): the output directory to save protobuf files.162        """163        logger = logging.getLogger(__name__)164        logger.info("Saving model to {} ...".format(output_dir))165        if not PathManager.exists(output_dir):166            PathManager.mkdirs(output_dir)167 168        with PathManager.open(os.path.join(output_dir, "model.pb"), "wb") as f:169            f.write(self._predict_net.SerializeToString())170        with PathManager.open(os.path.join(output_dir, "model.pbtxt"), "w") as f:171            f.write(str(self._predict_net))172        with PathManager.open(os.path.join(output_dir, "model_init.pb"), "wb") as f:173            f.write(self._init_net.SerializeToString())174 175    def save_graph(self, output_file, inputs=None):176        """177        Save the graph as SVG format.178 179        Args:180            output_file (str): a SVG file181            inputs: optional inputs given to the model.182                If given, the inputs will be used to run the graph to record183                shape of every tensor. The shape information will be184                saved together with the graph.185        """186        from .caffe2_export import run_and_save_graph187 188        if inputs is None:189            save_graph(self._predict_net, output_file, op_only=False)190        else:191            size_divisibility = get_pb_arg_vali(self._predict_net, "size_divisibility", 0)192            device = get_pb_arg_vals(self._predict_net, "device", b"cpu").decode("ascii")193            inputs = convert_batched_inputs_to_c2_format(inputs, size_divisibility, device)194            inputs = [x.cpu().numpy() for x in inputs]195            run_and_save_graph(self._predict_net, self._init_net, inputs, output_file)196 197    @staticmethod198    def load_protobuf(dir):199        """200        Args:201            dir (str): a directory used to save Caffe2Model with202                :meth:`save_protobuf`.203                The files "model.pb" and "model_init.pb" are needed.204 205        Returns:206            Caffe2Model: the caffe2 model loaded from this directory.207        """208        predict_net = caffe2_pb2.NetDef()209        with PathManager.open(os.path.join(dir, "model.pb"), "rb") as f:210            predict_net.ParseFromString(f.read())211 212        init_net = caffe2_pb2.NetDef()213        with PathManager.open(os.path.join(dir, "model_init.pb"), "rb") as f:214            init_net.ParseFromString(f.read())215 216        return Caffe2Model(predict_net, init_net)217 218    def __call__(self, inputs):219        """220        An interface that wraps around a Caffe2 model and mimics detectron2's models'221        input/output format. See details about the format at :doc:`/tutorials/models`.222        This is used to compare the outputs of caffe2 model with its original torch model.223 224        Due to the extra conversion between Pytorch/Caffe2, this method is not meant for225        benchmark. Because of the conversion, this method also has dependency226        on detectron2 in order to convert to detectron2's output format.227        """228        if self._predictor is None:229            self._predictor = ProtobufDetectionModel(self._predict_net, self._init_net)230        return self._predictor(inputs)231