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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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 |
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())