Python caffe2_pb2.NetDef方法代码示例

# 需要导入模块: from caffe2.proto import caffe2_pb2 [as 别名]
# 或者: from caffe2.proto.caffe2_pb2 import NetDef [as 别名]
def load_model_pb(net_file, init_file=None, is_run_init=True, is_create_net=True):
    net = core.Net("net")
    if net_file is not None:
        net.Proto().ParseFromString(open(net_file, "rb").read())

    if init_file is None:
        fn, ext = os.path.splitext(net_file)
        init_file = fn + "_init" + ext

    init_net = caffe2_pb2.NetDef()
    init_net.ParseFromString(open(init_file, "rb").read())

    if is_run_init:
        workspace.RunNetOnce(init_net)
        create_blobs_if_not_existed(net.external_inputs)
        if net.Proto().name == "":
            net.Proto().name = "net"
    if is_create_net:
        workspace.CreateNet(net)

    return (net, init_net) 

# 需要导入模块: from caffe2.proto import caffe2_pb2 [as 别名]
# 或者: from caffe2.proto.caffe2_pb2 import NetDef [as 别名]
def gen_init_net_from_blobs(blobs, blobs_to_use=None, excluded_blobs=None):
    ''' Generate an initialization net based on a blob dict '''
    ret = caffe2_pb2.NetDef()
    if blobs_to_use is None:
        blobs_to_use = {x for x in blobs}
    else:
        blobs_to_use = copy.deepcopy(blobs_to_use)
    if excluded_blobs is not None:
        blobs_to_use = [x for x in blobs_to_use if x not in excluded_blobs]
    for name in blobs_to_use:
        blob = blobs[name]
        if isinstance(blob, str):
            print('Blob {} with type {} is not supported in generating init net,'
                  ' skipped.'.format(name, type(blob)))
            continue
        add_tensor(ret, name, blob)

    return ret 

# 需要导入模块: from caffe2.proto import caffe2_pb2 [as 别名]
# 或者: from caffe2.proto.caffe2_pb2 import NetDef [as 别名]
def _propagate_device_option(net_def):
    '''
    Propagate the device options from net to operators.

    Args:
        net_def: A caffe2_pb2.NetDef representing a computation graph. The graph
            consists of Caffe2 operators.

    Returns:
        None. Iterates through all ops contained within the net. For each op,
            modifies the op device_option in-place to be the net device_option
            if the op has no pre-existing device_option, and leaves the op as-is
            if it already has a device_option.
    '''
    if not net_def.HasField("device_option"):
        return
    for op in net_def.op:
        if not op.HasField("device_option"):
            op.device_option.CopyFrom(net_def.device_option) 

# 需要导入模块: from caffe2.proto import caffe2_pb2 [as 别名]
# 或者: from caffe2.proto.caffe2_pb2 import NetDef [as 别名]
def nets_to_graph_def(nets, shapes=None, **kwargs):
    '''
    Convert a set of Caffe2 nets to a Tensorflow graph.

    Args:
        nets: List of core.Nets. core.Net is a wrapper around a NetDef protobuf.
            The corresponding protobuf can be extracted using .Proto().
        shapes: Dictionary mapping blob names to their shapes/dimensions.

    Returns:
        Call to protos_to_graph_def() with the extracted NetDef protobufs and
            **kwargs. See _operators_to_graph_def for detailed **kwargs.
    '''
    # if shapes is None:
    #     shapes = _try_get_shapes(nets)
    # _try_get_shapes(nets) depends on workspace.InferShapesAndTypes(nets),
    # which is currently broken (segfault). We omit the shapes for now.
    shapes = {}
    nets = [copy.deepcopy(net.Proto()) for net in nets]
    shapes = copy.deepcopy(shapes)
    return protos_to_graph_def(nets, shapes, **kwargs) 

# 需要导入模块: from caffe2.proto import caffe2_pb2 [as 别名]
# 或者: from caffe2.proto.caffe2_pb2 import NetDef [as 别名]
def protos_to_graph_def(net_defs, shapes=None, **kwargs):
    '''
    Convert a set of Caffe2 net definitions to a Tensorflow graph.

    Args:
        net_defs: List of caffe2_pb2.NetDef protobufs representing computation
            graphs.
        shapes: Dictionary mapping blob names to their shapes/dimensions.

    Returns:
        Call to _operators_to_graph_def() with the extracted operators from the
            NetDefs and **kwargs. See _operators_to_graph_def for detailed
            **kwargs.
    '''
    for net in net_defs:
        _propagate_device_option(net)
    shapes = copy.deepcopy(shapes or {})
    ops = [op for net_def in net_defs for op in net_def.op]
    return _operators_to_graph_def(shapes, ops, **kwargs) 

# 需要导入模块: from caffe2.proto import caffe2_pb2 [as 别名]
# 或者: from caffe2.proto.caffe2_pb2 import NetDef [as 别名]
def ConvertTensorProtosToInitNet(net_params, input_name):
    """Takes the net_params returned from TranslateModel, and wrap it as an
    init net that contain GivenTensorFill.

    This is a very simple feature that only works with float tensors, and is
    only intended to be used in an environment where you want a single
    initialization file - for more complex cases, use a db to store the
    parameters.
    """
    init_net = caffe2_pb2.NetDef()
    for tensor in net_params.protos:
        if len(tensor.float_data) == 0:
            raise RuntimeError(
                "Only float tensors are supported in this util.")
        op = core.CreateOperator(
            "GivenTensorFill", [], [tensor.name],
            arg=[
                utils.MakeArgument("shape", list(tensor.dims)),
                utils.MakeArgument("values", tensor.float_data)])
        init_net.op.extend([op])
    init_net.op.extend([core.CreateOperator("ConstantFill", [], [input_name], shape=[1])])
    return init_net 

# 需要导入模块: from caffe2.proto import caffe2_pb2 [as 别名]
# 或者: from caffe2.proto.caffe2_pb2 import NetDef [as 别名]
def __init__(self, predict_net, init_net):
        logger.info("Initializing ProtobufModel ...")
        super().__init__()
        assert isinstance(predict_net, caffe2_pb2.NetDef)
        assert isinstance(init_net, caffe2_pb2.NetDef)
        self.ws_name = "__ws_tmp__"
        self.net = core.Net(predict_net)

        with ScopedWS(self.ws_name, is_reset=True, is_cleanup=False) as ws:
            ws.RunNetOnce(init_net)
            for blob in self.net.Proto().external_input:
                if blob not in ws.Blobs():
                    ws.CreateBlob(blob)
            ws.CreateNet(self.net)

        self._error_msgs = set() 

# 需要导入模块: from caffe2.proto import caffe2_pb2 [as 别名]
# 或者: from caffe2.proto.caffe2_pb2 import NetDef [as 别名]
def construct_init_net_from_params(
    params: Dict[str, Any], device_options: Optional[Dict[str, caffe2_pb2.DeviceOption]] = None
) -> caffe2_pb2.NetDef:
    """
    Construct the init_net from params dictionary
    """
    init_net = caffe2_pb2.NetDef()
    device_options = device_options or {}
    for name, blob in params.items():
        if isinstance(blob, str):
            logger.warning(
                (
                    "Blob {} with type {} is not supported in generating init net,"
                    " skipped.".format(name, type(blob))
                )
            )
            continue
        init_net.op.extend(
            [create_const_fill_op(name, blob, device_option=device_options.get(name, None))]
        )
        init_net.external_output.append(name)
    return init_net 

# 需要导入模块: from caffe2.proto import caffe2_pb2 [as 别名]
# 或者: from caffe2.proto.caffe2_pb2 import NetDef [as 别名]
def rename_op_output(predict_net: caffe2_pb2.NetDef, op_id: int, output_id: int, new_name: str):
    """
    Rename the op_id-th operator in predict_net, change it's output_id-th input's
        name to the new_name. It also does automatic re-route and change
        external_output and if necessary.
    - It allows multiple consumers of its output.
    - This function modifies predict_net in-place, doesn't need init_net.
    """
    assert isinstance(predict_net, caffe2_pb2.NetDef)

    ssa, blob_versions = core.get_ssa(predict_net)

    versioned_inputs, versioned_outputs = ssa[op_id]
    old_name, version = versioned_outputs[output_id]

    # update predict_net
    _rename_versioned_blob_in_proto(
        predict_net, old_name, new_name, version, ssa, {}, blob_versions
    ) 

# 需要导入模块: from caffe2.proto import caffe2_pb2 [as 别名]
# 或者: from caffe2.proto.caffe2_pb2 import NetDef [as 别名]
def identify_reshape_sub_graph(predict_net: caffe2_pb2.NetDef) -> List[List[int]]:
    """
    Idenfity the reshape sub-graph in a protobuf.
    The reshape sub-graph is defined as matching the following pattern:

    (input_blob) -> Op_1 -> ... -> Op_N -> (new_shape) -─┐
        └-------------------------------------------> Reshape -> (output_blob)

    Return:
        List of sub-graphs, each sub-graph is represented as a list of indices
        of the relavent ops, [Op_1, Op_2, ..., Op_N, Reshape]
    """

    ssa, _ = core.get_ssa(predict_net)

    ret = []
    for i, op in enumerate(predict_net.op):
        if op.type == "Reshape":
            assert len(op.input) == 2
            input_ssa = ssa[i][0]
            data_source = input_ssa[0]
            shape_source = input_ssa[1]
            op_indices = _get_dependency_chain(ssa, shape_source, data_source)
            ret.append(op_indices + [i])
    return ret 

# 需要导入模块: from caffe2.proto import caffe2_pb2 [as 别名]
# 或者: from caffe2.proto.caffe2_pb2 import NetDef [as 别名]
def remove_dead_end_ops(net_def: caffe2_pb2.NetDef):
    """ remove ops if its output is not used or not in external_output """
    ssa, versions = core.get_ssa(net_def)
    versioned_external_output = [(name, versions[name]) for name in net_def.external_output]
    consumer_map = get_consumer_map(ssa)
    removed_op_ids = set()

    def _is_dead_end(versioned_blob):
        return not (
            versioned_blob in versioned_external_output
            or (
                len(consumer_map[versioned_blob]) > 0
                and all(x[0] not in removed_op_ids for x in consumer_map[versioned_blob])
            )
        )

    for i, ssa_i in reversed(list(enumerate(ssa))):
        versioned_outputs = ssa_i[1]
        if all(_is_dead_end(outp) for outp in versioned_outputs):
            removed_op_ids.add(i)

    # simply removing those deadend ops should have no effect to external_output
    new_ops = [op for i, op in enumerate(net_def.op) if i not in removed_op_ids]
    del net_def.op[:]
    net_def.op.extend(new_ops) 

# 需要导入模块: from caffe2.proto import caffe2_pb2 [as 别名]
# 或者: from caffe2.proto.caffe2_pb2 import NetDef [as 别名]
def caffe2_to_onnx(caffe2_model_name, caffe2_model_dir):
    caffe2_init_proto = caffe2_pb2.NetDef()
    caffe2_predict_proto = caffe2_pb2.NetDef()

    with open(os.path.join(caffe2_model_dir, 'init_net.pb'), 'rb') as f:
        caffe2_init_proto.ParseFromString(f.read())
        caffe2_init_proto.name = '{}_init'.format(caffe2_model_name)
    with open(os.path.join(caffe2_model_dir, 'predict_net.pb'), 'rb') as f:
        caffe2_predict_proto.ParseFromString(f.read())
        caffe2_predict_proto.name = caffe2_model_name
    with open(os.path.join(caffe2_model_dir, 'value_info.json'), 'rb') as f:
        value_info = json.loads(f.read())

    print('Converting Caffe2 model {} in {} to ONNX format'.format(caffe2_model_name, caffe2_model_dir))
    onnx_model = caffe2.python.onnx.frontend.caffe2_net_to_onnx_model(
        init_net=caffe2_init_proto,
        predict_net=caffe2_predict_proto,
        value_info=value_info
    )

    return onnx_model, caffe2_init_proto, caffe2_predict_proto 

# 需要导入模块: from caffe2.proto import caffe2_pb2 [as 别名]
# 或者: from caffe2.proto.caffe2_pb2 import NetDef [as 别名]
def write_graph(self, model_or_nets_or_protos=None, **kwargs):
        '''Write graph to the summary.'''
        if isinstance(model_or_nets_or_protos, cnn.CNNModelHelper):
            current_graph, track_blob_names = model_to_graph(model_or_nets_or_protos, **kwargs)
        elif isinstance(model_or_nets_or_protos, list):
            if isinstance(model_or_nets_or_protos[0], core.Net):
                current_graph, track_blob_names = nets_to_graph(model_or_nets_or_protos, **kwargs)
            elif isinstance(model_or_nets_or_protos[0], caffe2_pb2.NetDef):
                current_graph, track_blob_names = protos_to_graph(model_or_nets_or_protos, **kwargs)
            else:
                raise NotImplementedError
        else:
            raise NotImplementedError
        self._file_writer.add_graph(current_graph)
        self._track_blob_names = track_blob_names
        # Once the graph is built, one can just map the blobs
        self.check_names()
        self.sort_out_names() 

# 需要导入模块: from caffe2.proto import caffe2_pb2 [as 别名]
# 或者: from caffe2.proto.caffe2_pb2 import NetDef [as 别名]
def main():
    args = make_args()
    config = configparser.ConfigParser()
    utils.load_config(config, args.config)
    for cmd in args.modify:
        utils.modify_config(config, cmd)
    with open(os.path.expanduser(os.path.expandvars(args.logging)), 'r') as f:
        logging.config.dictConfig(yaml.load(f))
    torch.manual_seed(args.seed)
    model_dir = utils.get_model_dir(config)
    init_net = caffe2_pb2.NetDef()
    with open(os.path.join(model_dir, 'init_net.pb'), 'rb') as f:
        init_net.ParseFromString(f.read())
    predict_net = caffe2_pb2.NetDef()
    with open(os.path.join(model_dir, 'predict_net.pb'), 'rb') as f:
        predict_net.ParseFromString(f.read())
    p = workspace.Predictor(init_net, predict_net)
    height, width = tuple(map(int, config.get('image', 'size').split()))
    tensor = torch.randn(1, 3, height, width)
    # Checksum
    output = p.run([tensor.numpy()])
    for key, a in [
        ('tensor', tensor.cpu().numpy()),
        ('output', output[0]),
    ]:
        print('\t'.join(map(str, [key, a.shape, utils.abs_mean(a), hashlib.md5(a.tostring()).hexdigest()]))) 

# 需要导入模块: from caffe2.proto import caffe2_pb2 [as 别名]
# 或者: from caffe2.proto.caffe2_pb2 import NetDef [as 别名]
def main():
    args = parser.parse_args()
    args.gpu_id = 0

    model = model_helper.ModelHelper(name="le_net", init_params=False)

    # Bring in the init net from init_net.pb
    init_net_proto = caffe2_pb2.NetDef()
    with open(args.c2_init, "rb") as f:
        init_net_proto.ParseFromString(f.read())
    model.param_init_net = core.Net(init_net_proto)  # model.param_init_net.AppendNet(core.Net(init_net_proto)) #

    # bring in the predict net from predict_net.pb
    predict_net_proto = caffe2_pb2.NetDef()
    with open(args.c2_predict, "rb") as f:
        predict_net_proto.ParseFromString(f.read())
    model.net = core.Net(predict_net_proto)  # model.net.AppendNet(core.Net(predict_net_proto))

    # CUDA performance not impressive
    #device_opts = core.DeviceOption(caffe2_pb2.PROTO_CUDA, args.gpu_id)
    #model.net.RunAllOnGPU(gpu_id=args.gpu_id, use_cudnn=True)
    #model.param_init_net.RunAllOnGPU(gpu_id=args.gpu_id, use_cudnn=True)

    input_blob = model.net.external_inputs[0]
    model.param_init_net.GaussianFill(
        [],
        input_blob.GetUnscopedName(),
        shape=(args.batch_size, 3, args.img_size, args.img_size),
        mean=0.0,
        std=1.0)
    workspace.RunNetOnce(model.param_init_net)
    workspace.CreateNet(model.net, overwrite=True)
    workspace.BenchmarkNet(model.net.Proto().name, 5, 20, True)