Python caffe2.proto.caffe2_pb2.NetDef() Examples

The following are 30 code examples of caffe2.proto.caffe2_pb2.NetDef(). You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may also want to check out all available functions/classes of the module caffe2.proto.caffe2_pb2 , or try the search function .
Example #1
Source File: model_convert_utils.py    From Detectron-Cascade-RCNN with Apache License 2.0 6 votes vote down vote up
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 
Example #2
Source File: model_utils.py    From inference with Apache License 2.0 6 votes vote down vote up
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) 
Example #3
Source File: model_convert_utils.py    From Clustered-Object-Detection-in-Aerial-Image with Apache License 2.0 6 votes vote down vote up
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 
Example #4
Source File: model_convert_utils.py    From seg_every_thing with Apache License 2.0 6 votes vote down vote up
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 
Example #5
Source File: caffe2_inference.py    From detectron2 with Apache License 2.0 6 votes vote down vote up
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() 
Example #6
Source File: shared.py    From detectron2 with Apache License 2.0 6 votes vote down vote up
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 
Example #7
Source File: shared.py    From detectron2 with Apache License 2.0 6 votes vote down vote up
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
    ) 
Example #8
Source File: shared.py    From detectron2 with Apache License 2.0 6 votes vote down vote up
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 
Example #9
Source File: shared.py    From detectron2 with Apache License 2.0 6 votes vote down vote up
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) 
Example #10
Source File: model_utils.py    From models with Apache License 2.0 6 votes vote down vote up
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) 
Example #11
Source File: model_convert_utils.py    From Detectron with Apache License 2.0 6 votes vote down vote up
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 
Example #12
Source File: model_convert_utils.py    From Detectron-DA-Faster-RCNN with Apache License 2.0 6 votes vote down vote up
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 
Example #13
Source File: model_convert_utils.py    From CBNet with Apache License 2.0 6 votes vote down vote up
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 
Example #14
Source File: update-models-from-caffe2.py    From onnx-fb-universe with MIT License 6 votes vote down vote up
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 
Example #15
Source File: model_convert_utils.py    From NucleiDetectron with Apache License 2.0 6 votes vote down vote up
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 
Example #16
Source File: writer.py    From c2board with MIT License 6 votes vote down vote up
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() 
Example #17
Source File: model_convert_utils.py    From DetectAndTrack with Apache License 2.0 6 votes vote down vote up
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 
Example #18
Source File: caffe2_graph.py    From tensorboardX with MIT License 6 votes vote down vote up
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) 
Example #19
Source File: caffe_translator.py    From optimized-models with Apache License 2.0 6 votes vote down vote up
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 
Example #20
Source File: caffe2_graph.py    From tensorboardX with MIT License 6 votes vote down vote up
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) 
Example #21
Source File: caffe2_graph.py    From tensorboardX with MIT License 6 votes vote down vote up
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) 
Example #22
Source File: caffe2_graph.py    From tensorboardX with MIT License 6 votes vote down vote up
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) 
Example #23
Source File: caffe2_graph.py    From tensorboardX with MIT License 6 votes vote down vote up
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) 
Example #24
Source File: model_convert_utils.py    From KL-Loss with Apache License 2.0 6 votes vote down vote up
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 
Example #25
Source File: caffe2_benchmark.py    From gen-efficientnet-pytorch with Apache License 2.0 5 votes vote down vote up
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) 
Example #26
Source File: classification_no_db_example.py    From peters-stuff with GNU General Public License v3.0 5 votes vote down vote up
def load_net(INIT_NET, PREDICT_NET, device_opts):

    init_def = caffe2_pb2.NetDef()
    with open(INIT_NET, 'r') as f:
        init_def.ParseFromString(f.read())
        init_def.device_option.CopyFrom(device_opts)
        workspace.RunNetOnce(init_def.SerializeToString())

    net_def = caffe2_pb2.NetDef()
    with open(PREDICT_NET, 'r') as f:
        net_def.ParseFromString(f.read())
        net_def.device_option.CopyFrom(device_opts)
        workspace.CreateNet(net_def.SerializeToString(), overwrite=True) 
Example #27
Source File: classification_no_db_example.py    From peters-stuff with GNU General Public License v3.0 5 votes vote down vote up
def save_net(INIT_NET, PREDICT_NET, model) :

    with open(PREDICT_NET, 'wb') as f:
        f.write(model.net._net.SerializeToString())
    init_net = caffe2_pb2.NetDef()
    for param in model.params:
        blob = workspace.FetchBlob(param)
        shape = blob.shape
        op = core.CreateOperator("GivenTensorFill", [], [param],arg=[ utils.MakeArgument("shape", shape),utils.MakeArgument("values", blob)])
        init_net.op.extend([op])
    init_net.op.extend([core.CreateOperator("ConstantFill", [], ["data"], shape=(1,30,30))])
    with open(INIT_NET, 'wb') as f:
        f.write(init_net.SerializeToString()) 
Example #28
Source File: segmentation_no_db_example.py    From peters-stuff with GNU General Public License v3.0 5 votes vote down vote up
def load_net(INIT_NET, PREDICT_NET, device_opts):

    init_def = caffe2_pb2.NetDef()
    with open(INIT_NET, 'r') as f:
        init_def.ParseFromString(f.read())
        init_def.device_option.CopyFrom(device_opts)
        workspace.RunNetOnce(init_def.SerializeToString())

    net_def = caffe2_pb2.NetDef()
    with open(PREDICT_NET, 'r') as f:
        net_def.ParseFromString(f.read())
        net_def.device_option.CopyFrom(device_opts)
        workspace.CreateNet(net_def.SerializeToString(), overwrite=True) 
Example #29
Source File: segmentation_no_db_example.py    From peters-stuff with GNU General Public License v3.0 5 votes vote down vote up
def save_net(INIT_NET, PREDICT_NET, model) :

    with open(PREDICT_NET, 'wb') as f:
        f.write(model.net._net.SerializeToString())
    init_net = caffe2_pb2.NetDef()
    for param in model.params:
        #print param
        blob = workspace.FetchBlob(param)
        shape = blob.shape
        op = core.CreateOperator("GivenTensorFill", [], [param],arg=[ utils.MakeArgument("shape", shape),utils.MakeArgument("values", blob)])
        init_net.op.extend([op])
    init_net.op.extend([core.CreateOperator("ConstantFill", [], ["data"], shape=get_data(1)[0][0,:,:,:].shape)])
    with open(INIT_NET, 'wb') as f:
        f.write(init_net.SerializeToString()) 
Example #30
Source File: estimate.py    From openpose-pytorch with GNU Lesser General Public License v3.0 5 votes vote down vote up
def __init__(self, args, config):
        self.args = args
        self.config = config
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.cache_dir = utils.get_cache_dir(config)
        self.model_dir = utils.get_model_dir(config)
        _, self.num_parts = utils.get_dataset_mappers(config)
        self.limbs_index = utils.get_limbs_index(config)
        if args.debug is None:
            self.draw_cluster = utils.visualize.DrawCluster(colors=args.colors, thickness=args.thickness)
        else:
            self.draw_feature = utils.visualize.DrawFeature()
            s = re.search('(-?[0-9]+)([a-z]+)(-?[0-9]+)', args.debug)
            stage = int(s.group(1))
            name = s.group(2)
            channel = int(s.group(3))
            self.get_feature = lambda outputs: outputs[stage][name][0][channel]
        self.height, self.width = tuple(map(int, config.get('image', 'size').split()))
        if args.caffe:
            init_net = caffe2_pb2.NetDef()
            with open(os.path.join(self.model_dir, 'init_net.pb'), 'rb') as f:
                init_net.ParseFromString(f.read())
            predict_net = caffe2_pb2.NetDef()
            with open(os.path.join(self.model_dir, 'predict_net.pb'), 'rb') as f:
                predict_net.ParseFromString(f.read())
            p = workspace.Predictor(init_net, predict_net)
            self.inference = lambda tensor: [{'parts': torch.from_numpy(parts), 'limbs': torch.from_numpy(limbs)} for parts, limbs in zip(*[iter(p.run([tensor.detach().cpu().numpy()]))] * 2)]
        else:
            self.step, self.epoch, self.dnn, self.stages = self.load()
            self.inference = model.Inference(config, self.dnn, self.stages)
            self.inference.eval()
            if torch.cuda.is_available():
                self.inference.cuda()
            logging.info(humanize.naturalsize(sum(var.cpu().numpy().nbytes for var in self.inference.state_dict().values())))
        self.cap = self.create_cap()
        self.keys = set(args.keys)
        self.resize = transform.parse_transform(config, config.get('transform', 'resize_test'))
        self.transform_image = transform.get_transform(config, config.get('transform', 'image_test').split())
        self.transform_tensor = transform.get_transform(config, config.get('transform', 'tensor').split())