Class Caffe2Model

detectron2/export/api.py:187–298

A wrapper around the traced model in caffe2's pb format. Examples: :: model = Caffe2Model.load_protobuf("dir/with/pb/files") inputs = [{"image": img_tensor_CHW}] outputs = model(inputs)

class Caffe2Model(nn.Module):
    """
    A wrapper around the traced model in caffe2's pb format.

    Examples:
    ::
        model = Caffe2Model.load_protobuf("dir/with/pb/files")
        inputs = [{"image": img_tensor_CHW}]
        outputs = model(inputs)

    """

    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

    __init__.__HIDE_SPHINX_DOC__ = True

    @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

    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))
        if not PathManager.exists(output_dir):
            PathManager.mkdirs(output_dir)

        with PathManager.open(os.path.join(output_dir, "model.pb"), "wb") as f:
            f.write(self._predict_net.SerializeToString())
        with PathManager.open(os.path.join(output_dir, "model.pbtxt"), "w") as f:
            f.write(str(self._predict_net))
        with PathManager.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):
        """