Wrapper of a caffe2's protobuf model. It works just like nn.Module, but running caffe2 under the hood. Input/Output are Dict[str, tensor] whose keys are in external_input/output.
| 15 | |
| 16 | # ===== ref: mobile-vision's 'Caffe2Wrapper' class ====== |
| 17 | class 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 Dict[str, tensor] whose keys are in 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 instance |
| 32 | 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_input |
| 47 | |
| 48 | def _infer_output_devices(self, inputs): |
| 49 | """ |
| 50 | Returns: |
| 51 | list[str]: list of device for each external output |
| 52 | """ |
| 53 | |
| 54 | def _get_device_type(torch_tensor): |
| 55 | assert torch_tensor.device.type in ["cpu", "cuda"] |
| 56 | assert torch_tensor.device.index == 0 |
| 57 | return torch_tensor.device.type |
| 58 | |
| 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_devices |
| 70 | |
| 71 | def forward(self, inputs): |
| 72 | """ |
| 73 | Args: |
| 74 | inputs (tuple[torch.Tensor]) |