Support caffe2 (#6)

* feat(caffe2):registry caffe2

* feat(caffe2):support caffe2

source/python/neuropod/backends/caffe2/executor.py

@@ -0,0 +1,107 @@
+# Copyright (c) 2020 UATC, LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import os
+import json
+from six import string_types
+import caffe2
+from caffe2.python import workspace
+
+from neuropod.backends.neuropod_executor import NeuropodExecutor
+
+
+class Caffe2NeuropodExecutor(NeuropodExecutor):
+    """
+    Executes a Caffe2 neuropod
+    """
+    def __init__(self, neuropod_path):
+        """
+        Load a Caffe2 neuropod
+
+        :param  neuropod_path:      The path to a Caffe2 neuropod package
+        """
+        super(Caffe2NeuropodExecutor, self).__init__(neuropod_path)
+
+        neuropod_data_path = os.path.join(neuropod_path, "0", "data")
+
+        # Add the model to the neuropod
+        init_path = os.path.join(neuropod_data_path, "init_net.pb")
+        predict_path = os.path.join(neuropod_data_path, "predict_net.pb")
+
+        workspace.ResetWorkspace()
+        with open(init_path, "rb") as f:
+            init_net = f.read()
+        with open(predict_path, "rb") as f:
+            predict_net = f.read()
+        workspace.RunNetOnce(init_net)
+        workspace.CreateNet(predict_net)
+
+        self.model = workspace.Predictor(init_net, predict_net)
+        with open(os.path.join(neuropod_path, "0", "config.json"),
+                  "r") as config_file:
+            model_config = json.load(config_file)
+
+            # Get the node name mapping and store it
+            self.node_name_mapping = model_config["node_name_mapping"]
+
+    def forward(self, inputs):
+        """
+        Run inference using the specifed inputs.
+
+        :param  inputs:     A dict mapping input names to values. This must match the input
+                            spec in the neuropod config for the loaded model.
+                            Ex: {'x1': np.array([5]), 'x2': np.array([6])}
+                            *Note:* all the keys in this dict must be strings and all the
+                            values must be numpy arrays
+
+        :returns:   A dict mapping output names to values. All the keys
+                    in this dict are strings and all the values are numpy arrays.
+        """
+
+        # Convert the inputs to torch tensors and move to the appropriate device
+
+        output_dict = {}
+        feed_dict = {}
+
+        # Get the output nodes
+        for node in self.neuropod_config["output_spec"]:
+            neuropod_name = node["name"]
+
+            # Get the graph node
+            caffe2_name = self.node_name_mapping[neuropod_name]
+
+            # Add it to the output name
+            output_dict[neuropod_name] = caffe2_name
+
+        # Get the input nodes
+        for node in self.neuropod_config["input_spec"]:
+            neuropod_name = node["name"]
+
+            if neuropod_name not in inputs:
+                continue
+
+            # Get the graph node
+            caffe2_name = self.node_name_mapping[neuropod_name]
+
+            # Add it to the feed_dict
+            feed_dict[caffe2_name] = inputs[neuropod_name]
+
+        # forward to feed blobs
+        self.model.run(feed_dict)
+
+        neuropod_out = {}
+        for k, v in output_dict.items():
+            neuropod_out[k] = workspace.FetchBlob(v)
+        return neuropod_out

source/python/neuropod/backends/caffe2/packager.py

@@ -0,0 +1,82 @@
+# Copyright (c) 2020 UATC, LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import json
+import shutil
+
+from neuropod.utils.packaging_utils import packager
+
+
+@packager(platform="caffe2")
+def create_caffe2_neuropod(neuropod_path, input_spec, output_spec,
+                           node_name_mapping, predict_net, init_net, **kwargs):
+    """
+    Packages a TorchScript model as a neuropod package.
+
+    {common_doc_pre}
+
+    :param  node_name_mapping:  Mapping from a neuropod input/output name to a blob in the net. 
+
+                                !!! note ""
+                                    ***Example***:
+                                    ```
+                                    {
+                                        "x": "inputA",
+                                        "y": "inputB",
+                                        "out": "output",
+                                    }
+                                    ```
+
+    :param  init_net:           path a protobuf that has all of the network weights. 
+
+    :param  predict_net:        path a protobuf that defines the network.
+    {common_doc_post}
+    """
+
+    # Create a folder to store the model
+    neuropod_data_path = os.path.join(neuropod_path, "0", "data")
+    os.makedirs(neuropod_data_path)
+
+    # Add the model to the neuropod
+    init_path = os.path.join(neuropod_data_path, "init_net.pb")
+    predict_path = os.path.join(neuropod_data_path, "predict_net.pb")
+
+    shutil.copyfile(init_net, init_path)
+    shutil.copyfile(predict_net, predict_path)
+
+    # Make sure we have mappings for everything in the spec
+    expected_keys = set()
+    for spec in [input_spec, output_spec]:
+        for tensor in spec:
+            expected_keys.add(tensor["name"])
+
+    actual_keys = set(node_name_mapping.keys())
+    missing_keys = expected_keys - actual_keys
+
+    if len(missing_keys) > 0:
+        raise ValueError(
+            "Expected an item in `node_name_mapping` for every tensor in input_spec and output_spec. Missing: `{}`"
+            .format(missing_keys))
+
+    # We also need to save the node name mapping so we know how to run the model
+    # This is tensorflow specific config so it's not saved in the overall neuropod config
+    with open(os.path.join(neuropod_path, "0", "config.json"),
+              "w") as config_file:
+        json.dump(
+            {
+                "node_name_mapping": node_name_mapping,
+            },
+            config_file,
+        )

source/python/neuropod/loader.py

@@ -230,6 +230,10 @@ def load_neuropod(neuropod_path, _always_use_native=True, **kwargs):
         from neuropod.backends.onnx.executor import OnnxNeuropodExecutor
 
         return OnnxNeuropodExecutor(neuropod_path, **kwargs)
+    elif platform == "caffe2":
+        from neuropod.backends.caffe2.executor import Caffe2NeuropodExecutor
+
+        return Caffe2NeuropodExecutor(neuropod_path, **kwargs)       
     else:
         raise ValueError(
             "Invalid platform found in neuropod config: {}".format(platform)

source/python/neuropod/packagers.py

@@ -46,6 +46,7 @@ def __getattr__(self, packaging_function):
             "torchscript",
             "caffe",
             "onnx",
+            "caffe2",
         ]:
             raise RuntimeError(
                 "Tried to get an invalid attribute on neuropod.packagers ({})".format(