Arulkumar03/Fox_Sheep_Detector_Computer_Vision_model
0
1# Copyright (c) Facebook, Inc. and its affiliates.2 3import logging4import numpy as np5from itertools import count6import torch7from caffe2.proto import caffe2_pb28from caffe2.python import core9 10from .caffe2_modeling import META_ARCH_CAFFE2_EXPORT_TYPE_MAP, convert_batched_inputs_to_c2_format11from .shared import ScopedWS, get_pb_arg_vali, get_pb_arg_vals, infer_device_type12 13logger = logging.getLogger(__name__)14 15 16# ===== ref: mobile-vision predictor's 'Caffe2Wrapper' class ======17class ProtobufModel(torch.nn.Module):18 """19 Wrapper of a caffe2's protobuf model.20 It works just like nn.Module, but running caffe2 under the hood.21 Input/Output are tuple[tensor] that match the caffe2 net's external_input/output.22 """23 24 _ids = count(0)25 26 def __init__(self, predict_net, init_net):27 logger.info(f"Initializing ProtobufModel for: {predict_net.name} ...")28 super().__init__()29 assert isinstance(predict_net, caffe2_pb2.NetDef)30 assert isinstance(init_net, caffe2_pb2.NetDef)31 # create unique temporary workspace for each instance32 self.ws_name = "__tmp_ProtobufModel_{}__".format(next(self._ids))33 self.net = core.Net(predict_net)34 35 logger.info("Running init_net once to fill the parameters ...")36 with ScopedWS(self.ws_name, is_reset=True, is_cleanup=False) as ws:37 ws.RunNetOnce(init_net)38 uninitialized_external_input = []39 for blob in self.net.Proto().external_input:40 if blob not in ws.Blobs():41 uninitialized_external_input.append(blob)42 ws.CreateBlob(blob)43 ws.CreateNet(self.net)44 45 self._error_msgs = set()46 self._input_blobs = uninitialized_external_input47 48 def _infer_output_devices(self, inputs):49 """50 Returns:51 list[str]: list of device for each external output52 """53 54 def _get_device_type(torch_tensor):55 assert torch_tensor.device.type in ["cpu", "cuda"]56 assert torch_tensor.device.index == 057 return torch_tensor.device.type58 59 predict_net = self.net.Proto()60 input_device_types = {61 (name, 0): _get_device_type(tensor) for name, tensor in zip(self._input_blobs, inputs)62 }63 device_type_map = infer_device_type(64 predict_net, known_status=input_device_types, device_name_style="pytorch"65 )66 ssa, versions = core.get_ssa(predict_net)67 versioned_outputs = [(name, versions[name]) for name in predict_net.external_output]68 output_devices = [device_type_map[outp] for outp in versioned_outputs]69 return output_devices70 71 def forward(self, inputs):72 """73 Args:74 inputs (tuple[torch.Tensor])75 76 Returns:77 tuple[torch.Tensor]78 """79 assert len(inputs) == len(self._input_blobs), (80 f"Length of inputs ({len(inputs)}) "81 f"doesn't match the required input blobs: {self._input_blobs}"82 )83 84 with ScopedWS(self.ws_name, is_reset=False, is_cleanup=False) as ws:85 for b, tensor in zip(self._input_blobs, inputs):86 ws.FeedBlob(b, tensor)87 88 try:89 ws.RunNet(self.net.Proto().name)90 except RuntimeError as e:91 if not str(e) in self._error_msgs:92 self._error_msgs.add(str(e))93 logger.warning("Encountered new RuntimeError: \n{}".format(str(e)))94 logger.warning("Catch the error and use partial results.")95 96 c2_outputs = [ws.FetchBlob(b) for b in self.net.Proto().external_output]97 # Remove outputs of current run, this is necessary in order to98 # prevent fetching the result from previous run if the model fails99 # in the middle.100 for b in self.net.Proto().external_output:101 # Needs to create uninitialized blob to make the net runable.102 # This is "equivalent" to: ws.RemoveBlob(b) then ws.CreateBlob(b),103 # but there'no such API.104 ws.FeedBlob(b, f"{b}, a C++ native class of type nullptr (uninitialized).")105 106 # Cast output to torch.Tensor on the desired device107 output_devices = (108 self._infer_output_devices(inputs)109 if any(t.device.type != "cpu" for t in inputs)110 else ["cpu" for _ in self.net.Proto().external_output]111 )112 113 outputs = []114 for name, c2_output, device in zip(115 self.net.Proto().external_output, c2_outputs, output_devices116 ):117 if not isinstance(c2_output, np.ndarray):118 raise RuntimeError(119 "Invalid output for blob {}, received: {}".format(name, c2_output)120 )121 outputs.append(torch.tensor(c2_output).to(device=device))122 return tuple(outputs)123 124 125class ProtobufDetectionModel(torch.nn.Module):126 """127 A class works just like a pytorch meta arch in terms of inference, but running128 caffe2 model under the hood.129 """130 131 def __init__(self, predict_net, init_net, *, convert_outputs=None):132 """133 Args:134 predict_net, init_net (core.Net): caffe2 nets135 convert_outptus (callable): a function that converts caffe2136 outputs to the same format of the original pytorch model.137 By default, use the one defined in the caffe2 meta_arch.138 """139 super().__init__()140 self.protobuf_model = ProtobufModel(predict_net, init_net)141 self.size_divisibility = get_pb_arg_vali(predict_net, "size_divisibility", 0)142 self.device = get_pb_arg_vals(predict_net, "device", b"cpu").decode("ascii")143 144 if convert_outputs is None:145 meta_arch = get_pb_arg_vals(predict_net, "meta_architecture", b"GeneralizedRCNN")146 meta_arch = META_ARCH_CAFFE2_EXPORT_TYPE_MAP[meta_arch.decode("ascii")]147 self._convert_outputs = meta_arch.get_outputs_converter(predict_net, init_net)148 else:149 self._convert_outputs = convert_outputs150 151 def _convert_inputs(self, batched_inputs):152 # currently all models convert inputs in the same way153 return convert_batched_inputs_to_c2_format(154 batched_inputs, self.size_divisibility, self.device155 )156 157 def forward(self, batched_inputs):158 c2_inputs = self._convert_inputs(batched_inputs)159 c2_results = self.protobuf_model(c2_inputs)160 c2_results = dict(zip(self.protobuf_model.net.Proto().external_output, c2_results))161 return self._convert_outputs(batched_inputs, c2_inputs, c2_results)162